Add info message, if model weights are loaded.

Author: Lucas-Steinmann

train.py

@@ -42,9 +42,9 @@ def create_parser():
     parser.add_argument("-c", "--config-file", type=str, required=False,
                         help="Path to a yaml configuration file.")
     parser.add_argument("--log", required=False, type=str, help="Path to the log file destination.")
-    parser.add_argument("--save_path", required=False, type=str, default="", 
+    parser.add_argument("--save_path", required=False, type=str, default="",
                         help="Path to the model destination. If empty, model won't be saved.")
-    parser.add_argument("--load_path", required=False, type=str, default="", 
+    parser.add_argument("--load_path", required=False, type=str, default="",
                         help="Path to the saved model. If empty, model won't be loaded.")
     return parser
 
@@ -181,9 +181,10 @@ def train(args: Namespace, seed: int = 0, verbose: bool = False) -> Tuple[List[D
     seq_length = 12
 
     if args.load_path is not None and Path(args.load_path).is_file():
-      model = torch.load(args.load_path)      
-    else: 
-      model = build_model(11, seq_length, args.batch_size).to(device)
+        print("Loading model weights from: " + args.load_path)
+        model = torch.load(args.load_path)
+    else:
+        model = build_model(11, seq_length, args.batch_size).to(device)
 
     optimizer = torch.optim.Adam(model.parameters(), lr=args.lr)
     floss = loss_func()
@@ -256,7 +257,7 @@ def train(args: Namespace, seed: int = 0, verbose: bool = False) -> Tuple[List[D
         write_to_csv(history_item, args.log, write_header=epoch == 0, append=epoch != 0)
 
         if args.save_path is not None:
-          torch.save(model, args.save_path)          
+            torch.save(model, args.save_path)
 
     # Test here
     test_results = test(model, dataloaders['test'], verbose)