From b0eaff36e6aefd45c0fe89bb2ab86a495e4e735f Mon Sep 17 00:00:00 2001 From: piero Date: Wed, 27 Nov 2019 21:43:43 -0800 Subject: [PATCH] Added a +1 to epoch when saving weights --- examples/run_pplm_discrim_train.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/examples/run_pplm_discrim_train.py b/examples/run_pplm_discrim_train.py index 5291ad4b51..fccfb14426 100644 --- 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__":