CUDA specific code again

allenai · Jan 2, 2025 · fb4466c · fb4466c
1 parent 51a9294
commit fb4466c
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/olmo_core/train/trainer.py b/src/olmo_core/train/trainer.py
@@ -1299,7 +1299,8 @@ def _fit_epoch(self):
 
             if first_batch or self.global_step % self.metrics_collect_interval == 0:
                 self._log_metrics()
-                torch.cuda.set_sync_debug_mode("warn")
+                if torch.cuda.is_available():
+                    torch.cuda.set_sync_debug_mode("warn")
 
             first_batch = False