Update sizer.py

mila-iqia · Jun 7, 2024 · 702c1c6 · 702c1c6
1 parent 20b1085
commit 702c1c6
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/milabench/sizer.py b/milabench/sizer.py
@@ -106,7 +106,12 @@ def auto_size(self, benchmark, capacity):
             return None
 
         config = self.benchscaling(benchmark)
-
+        model = config.get("model", None)
+
+        if model is None:
+            print(f"Missing batch-size model for {benchmark}")
+            return 1
+
         data = list(sorted(config["model"].items(), key=lambda x: x[0]))
         mem = [to_octet(v[1]) for v in data]
         size = [float(v[0]) for v in data]