Update train_net.py

IDEA-Research · Oct 24, 2023 · e6998ed · e6998ed
1 parent f95cd41
commit e6998ed
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/train_net.py b/train_net.py
@@ -443,7 +443,10 @@ def main(args=None):
     trainer.resume_or_load(resume=args.resume)
     if args.original_load:  # loading checkpoints with different name prefix
         print("using original loading")
-        trainer._trainer.model.module = trainer._trainer.model.module.from_pretrained(cfg.MODEL.WEIGHTS)
+        try:
+            trainer._trainer.model.module = trainer._trainer.model.module.from_pretrained(cfg.MODEL.WEIGHTS)
+        except Exception as e: # for debugging
+            trainer._trainer.model = trainer._trainer.model.from_pretrained(cfg.MODEL.WEIGHTS)
     return trainer.train()