rm cpu

layumi · Jan 24, 2024 · e03564e · e03564e
1 parent d81db6f
commit e03564e
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/utils.py b/utils.py
@@ -56,10 +56,10 @@ def save_network(network, dirname, epoch_label, local_rank=-1):
 
     if local_rank>-1:
         if local_rank == 0: # save the main process model
-            torch.save(network.cpu().state_dict(), save_path)
+            torch.save(network.state_dict(), save_path)
             network.cuda(local_rank)
     else:
-        torch.save(network.cpu().state_dict(), save_path)
+        torch.save(network.state_dict(), save_path)
         network.cuda()