[engine] More comments on model.train() and model.to() usage

bob/ip/binseg/engine/predictor.py

@@ -125,9 +125,10 @@ def run(model, data_loader, device, output_folder, overlayed_folder):
     os.makedirs(output_folder, exist_ok=True)
 
     logger.info(f"Device: {device}")
-    model.eval().to(device)
-    # Sigmoid for predictions
-    sigmoid = torch.nn.Sigmoid()
+
+    model.eval()  # set evaluation mode
+    model.to(device)  # set/cast parameters to device
+    sigmoid = torch.nn.Sigmoid()  # use sigmoid for predictions
 
     # Setup timers
     start_total_time = time.time()

bob/ip/binseg/engine/ssltrainer.py

@@ -300,7 +300,9 @@ def run(
         if arguments["epoch"] == 0:
             logwriter.writeheader()
 
-        model.train().to(device)
+        model.train()  #set training mode
+
+        model.to(device)  # set/cast parameters to device
         for state in optimizer.state.values():
             for k, v in state.items():
                 if isinstance(v, torch.Tensor):

bob/ip/binseg/engine/trainer.py

@@ -179,7 +179,9 @@ def run(
         if arguments["epoch"] == 0:
             logwriter.writeheader()
 
-        model.train().to(device)
+        model.train()  #set training mode
+
+        model.to(device)  # set/cast parameters to device
         for state in optimizer.state.values():
             for k, v in state.items():
                 if isinstance(v, torch.Tensor):