Update train.py

wonbeomjang · Feb 5, 2022 · 5bebb06 · 5bebb06
1 parent 76e2896
commit 5bebb06
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/train.py b/src/train.py
@@ -131,7 +131,7 @@ def quantize_model(self):
             return
 
         print('Load Best Model')
-        ckpt = f'{self.checkpoint_dir}/mobilenetv2.pt'
+        ckpt = f'{self.checkpoint_dir}/best.pt'
         save_info = torch.load(ckpt, map_location=self.device)
         self.net = save_info['model']
         self.net.load_state_dict(save_info['state_dict'])