fix save model bugs

2022-04-21 20:15:21 +08:00
parent b893316e41
commit 7ed12d218f
4 changed files with 9 additions and 9 deletions
@@ -52,7 +52,7 @@ class BaseModel(torch.nn.Module):
        torch.save(network.state_dict(), save_path)

    # helper saving function that can be used by subclasses
-    def save_network(self, network, network_label, epoch_label, gpu_ids):
+    def save_network(self, network, network_label, epoch_label, gpu_ids=None):
        save_filename = '%s_net_%s.pth' % (epoch_label, network_label)
        save_path = os.path.join(self.save_dir, save_filename)
        torch.save(network.cpu().state_dict(), save_path)