add truncation

prescient-design · Mar 4, 2025 · 98360cc · 98360cc
1 parent 763f4ab
commit 98360cc
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/lobster/data/_ume_haha_datamodule.py b/src/lobster/data/_ume_haha_datamodule.py
@@ -126,6 +126,7 @@ def _get_dataset_node(self, dataset_info: DatasetInfo, split: Split) -> Dataset:
             text_key="sequence",
             return_tensors="pt",
             padding="max_length",
+            truncation=True,
             max_length=self._tokenizer_max_length,
             num_workers=self._num_workers,
         )