Fix batch size issue

EleutherAI · Aug 18, 2023 · 96053f7 · 96053f7
1 parent 83f2c2a
commit 96053f7
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/inference.py b/inference.py
@@ -137,8 +137,7 @@ def get_batch_size(model_name: str) -> int:
  "6.9b": 64,
  "12b": 64,
  }
- model_size = ".".join(model_name.split(".")[1:])
- return size_batch_map[model_size]
+ return size_batch_map[model_name]
 
 
 def get_dataset(dataset_name: str, split_name: str, sample: int = None) -> pd.DataFrame: