Added a +1 to epoch when saving weights

2019-11-27 21:43:43 -08:00
parent 611961ade7
commit b0eaff36e6
1 changed files with 3 additions and 2 deletions
--- a/examples/run_pplm_discrim_train.py
+++ b/examples/run_pplm_discrim_train.py
@@ -545,10 +545,11 @@ def train_discriminator(
        if save_model:
            # torch.save(discriminator.state_dict(),
            #           "{}_discriminator_{}.pt".format(
-            #               args.dataset, epoch
+            #               args.dataset, epoch + 1
            #               ))
            torch.save(discriminator.get_classifier().state_dict(),
-                       "{}_classifier_head_epoch_{}.pt".format(dataset, epoch))
+                       "{}_classifier_head_epoch_{}.pt".format(dataset,
+                                                               epoch + 1))


 if __name__ == "__main__":