Merge pull request eriklindernoren#50 from chenyang1999/patch-3

updata the way about how to load pretrain model
0xzeny · Apr 26, 2019 · fd9f071 · fd9f071
2 parents d01fa6b + 22d8b45
commit fd9f071
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/implementations/esrgan/esrgan.py b/implementations/esrgan/esrgan.py
@@ -191,5 +191,5 @@
 
         if batches_done % opt.checkpoint_interval == 0:
             # Save model checkpoints
-            torch.save(generator.state_dict(), "saved_models/generator_%d.pth" % batches_done)
-            torch.save(discriminator.state_dict(), "saved_models/discriminator_%d.pth" % batches_done)
+            torch.save(generator.state_dict(), "saved_models/generator_%d.pth" % epoch)
+            torch.save(discriminator.state_dict(), "saved_models/discriminator_%d.pth" %epoch)