Skip to content

Commit c8f6153

Browse files
authored
fix cpu 'torch._C' has no attribute '_cuda_resetPeakMemoryStats' (#914)
1 parent 0d4b16c commit c8f6153

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

swift/trainers/mixin.py

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -516,7 +516,8 @@ def get_max_cuda_memory(self, device: Optional[Union[torch.device, int]] = None)
516516
mem = sum([float(mem) / 1024 / 1024 / 1024 for mem in mems])
517517
if self.max_memory < mem:
518518
self.max_memory = mem
519-
torch.cuda.reset_peak_memory_stats()
519+
if torch.cuda.is_available():
520+
torch.cuda.reset_peak_memory_stats()
520521
return mem
521522

522523
def _maybe_log_save_evaluate(self, tr_loss, *args, **kwargs):

0 commit comments

Comments
 (0)