always make mock batch on CPU

allenai · Jan 29, 2025 · 7ec9114 · 7ec9114
1 parent 7122b1d
commit 7ec9114
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/olmo_core/data/data_loader.py b/src/olmo_core/data/data_loader.py
@@ -434,7 +434,7 @@ def reshuffle(self, epoch: Optional[int] = None, in_memory: bool = False, **kwar
         self.build_and_save_global_indices(in_memory=in_memory)
 
     def get_mock_batch(self) -> Dict[str, Any]:
-        rng = torch.Generator(device=get_default_device())
+        rng = torch.Generator(device="cpu")
         rng.manual_seed(self.seed + self.dp_rank)
         num_instances = self.rank_batch_size // self.dataset.max_sequence_length
         input_ids = torch.randint(