Update train_test_utils/train.py

train_test_utils/train.py

@@ -1,14 +1,10 @@
 import torch
 import torch.nn as nn
 from torch.nn.utils import clip_grad_norm_
-from torch.utils.data import DataLoader
 from just_time_windows.Actor.actor import Actor
 
 
-
-
 def train_batch(actor, baseline, batch, optimizer, gradient_clipping=True, comparison_model=None, compute_cost_ratio=True):
-
     device = actor.device
 
     actor.train_mode()
@@ -16,7 +12,6 @@ def train_batch(actor, baseline, batch, optimizer, gradient_clipping=True, compa
     actor_output = actor(batch)
     actor_cost, log_probs = actor_output['total_time'], actor_output['log_probs']
 
-
     with torch.no_grad():
         baseline.greedy_search()
         baseline_output = baseline(batch)
@@ -29,19 +24,17 @@ def train_batch(actor, baseline, batch, optimizer, gradient_clipping=True, compa
 
     if gradient_clipping:
         for group in optimizer.param_groups:
-            clip_grad_norm_(
-                group['params'],
-                1,
-                norm_type=2
-            )
+            params = [p for p in group['params'] if p.grad is not None]
+            if params:
+                clip_grad_norm_(params, max_norm=1, norm_type=2)
 
     optimizer.step()
 
-    if compute_cost_ratio and (comparison_model is None):
-        normalize = actor.apply_normalization
-        comparison_model = Actor(model=None, num_neighbors_action=1, normalize=normalize, device=device)
-
     if compute_cost_ratio:
+        if comparison_model is None:
+            normalize = actor.apply_normalization
+            comparison_model = Actor(model=None, num_neighbors_action=1, normalize=normalize, device=device)
+
         with torch.no_grad():
             comp_output = comparison_model(batch)
             comp_cost = comp_output['total_time']
@@ -49,6 +42,5 @@ def train_batch(actor, baseline, batch, optimizer, gradient_clipping=True, compa
         a = comp_cost.sum().item()
         b = actor_cost.sum().item()
         return b / a
-    else:
-        return None
 
+    return None