Merge pull request #41 from jmercat/main

fix a small eval bug
state-spaces · Dec 9, 2023 · 165b07a · 165b07a
2 parents 0131c1e + 42f859d
commit 165b07a
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 None else 64
+        self._batch_size = batch_size if batch_size is not None else 64
         self._max_length = max_length
         self._device = torch.device(device)