From 0f22e7f7c63396e0a0839977ba8101f0d7bf1b04 Mon Sep 17 00:00:00 2001 From: Doyup Lee <33016192+LeeDoYup@users.noreply.github.com> Date: Sun, 31 Jan 2021 00:48:42 +0900 Subject: [PATCH] Fix bug in load_model in training close #12 --- models/fixmatch/fixmatch.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/models/fixmatch/fixmatch.py b/models/fixmatch/fixmatch.py index d5c0787f..5536bf92 100644 --- a/models/fixmatch/fixmatch.py +++ b/models/fixmatch/fixmatch.py @@ -266,6 +266,10 @@ def load_model(self, load_path): eval_model.load_state_dict(checkpoint[key]) elif key == 'it': self.it = checkpoint[key] + elif key == 'scheduler': + self.scheduler.load_state_dict(checkpoint[key]) + elif key == 'optimizer': + self.optimizer.load_state_dict(checkpoint[key]) else: getattr(self, key).load_state_dict(checkpoint[key]) self.print_fn(f"Check Point Loading: {key} is LOADED")