Fixed error caused by wrong parsing type in eval

state-spaces · Dec 9, 2023 · 3bb4491 · 3bb4491
1 parent 42f859d
commit 3bb4491
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/evals/lm_harness_eval.py b/evals/lm_harness_eval.py
@@ -23,7 +23,7 @@ def __init__(self, pretrained="state-spaces/mamba-2.8b", max_length=2048, batch_
         self.tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b")
         self.tokenizer.pad_token_id = self.tokenizer.eos_token_id
         self.vocab_size = self.tokenizer.vocab_size
-        self._batch_size = batch_size if batch_size is not None else 64
+        self._batch_size = int(batch_size) if batch_size is not None else 64
         self._max_length = max_length
         self._device = torch.device(device)