minor

plutonium-239 · May 22, 2024 · c14c4e6 · c14c4e6
1 parent 8853a5e
commit c14c4e6
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/experiments/paper_demo.py b/experiments/paper_demo.py
@@ -124,9 +124,9 @@
         )
 
         for model in models:
-            if model == "flan-t5":
+            if model in ["flan-t5", "memsave_flan-t5"]:
                 batch_size = 56
-            if model == "mistral-7b":
+            if model in ["mistral-7b", "memsave_mistral-7b"]:
                 batch_size = 16
             for estimate in estimators:
                 outputs = []

diff --git a/experiments/util/estimate.py b/experiments/util/estimate.py
@@ -297,9 +297,9 @@ def estimate_mem_savings(
             model_fn = models.transformer_model_fns.get(args.model)
             if args.model in models.hf_transformers_models:
                 model_fn_orig = model_fn
-                model_fn = lambda: models.TransformersModelWrapper(
+                model_fn = lambda: models.TransformersModelWrapper(  # noqa: E731
                     model_fn_orig, args.model
-                )  # noqa: E731
+                )
                 config = models.get_transformers_config(args.model)
                 # as per transformers.PretrainedConfig these 2 should be present in all models:
                 vocab_dim = config.vocab_size