int cast

mosaicml · Jun 13, 2024 · a8adcba · a8adcba
1 parent 4af9ab6
commit a8adcba
Showing 1 changed file with 1 addition and 5 deletions.
diff --git a/llmfoundry/data/finetuning/tasks.py b/llmfoundry/data/finetuning/tasks.py
@@ -680,11 +680,7 @@ def __getitem__(self, idx: int) -> Dict[str, Any]:
     def state_dict(self, num_samples: int,
                    from_beginning: bool) -> Dict[str, Any]:
         if self.packing_ratio is not None:
-            num_samples = self.packing_ratio * num_samples
-
-        print('+'*30)
-        print(num_samples)
-        print('+'*30)
+            num_samples = int(self.packing_ratio * num_samples)
 
         return super().state_dict(
             num_samples=num_samples,