Merge branch 'main' into release/2.0

modelscope · May 22, 2024 · ba3e277 · ba3e277
2 parents 3d98fef + c8f6153
commit ba3e277
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/swift/trainers/mixin.py b/swift/trainers/mixin.py
@@ -516,7 +516,8 @@ def get_max_cuda_memory(self, device: Optional[Union[torch.device, int]] = None)
         mem = sum([float(mem) / 1024 / 1024 / 1024 for mem in mems])
         if self.max_memory < mem:
             self.max_memory = mem
-        torch.cuda.reset_peak_memory_stats()
+        if torch.cuda.is_available():
+            torch.cuda.reset_peak_memory_stats()
         return mem
 
     def _maybe_log_save_evaluate(self, tr_loss, *args, **kwargs):