ソースを参照

解决断点续训失败问题

liyan 7 ヶ月 前
コミット
d053d4d95c
1 ファイル変更1 行追加1 行削除
  1. 1 1
      train.py

+ 1 - 1
train.py

@@ -341,7 +341,7 @@ def main(args):
         model_ema = utils.ExponentialMovingAverage(model_without_ddp, device=device, decay=1.0 - alpha)
 
     if args.resume:
-        checkpoint = torch.load(args.resume, map_location="cpu", weights_only=True)
+        checkpoint = torch.load(args.resume, map_location="cpu", weights_only=False)
         model_without_ddp.load_state_dict(checkpoint["model"])
         if not args.test_only:
             optimizer.load_state_dict(checkpoint["optimizer"])