Update train_test_utils/validation.py

train_test_utils/validation.py

@@ -1,24 +1,30 @@
-
 import torch
-import torch.nn as nn
-from torch.nn.utils import clip_grad_norm_
 from torch.utils.data import DataLoader
 
 
 def validation(actor, validation_dataset, batch_size):
-
+    """
+    Evaluate the actor model on the validation dataset.
+    
+    Args:
+        actor: Trained model to evaluate
+        validation_dataset: Dataset for validation
+        batch_size: Size of mini-batches used in evaluation
+
+    Returns:
+        Tensor of total costs for each sample in the validation set
+    """
+    actor.eval()  # Set model to evaluation mode
     val_dataloader = DataLoader(dataset=validation_dataset,
                                 batch_size=batch_size,
                                 collate_fn=validation_dataset.collate)
 
     scores = []
-    for batch in val_dataloader:
-        with torch.no_grad():
-            actor_output = actor(batch)
-            cost = actor_output['total_time']
 
-            scores.append(cost.reshape(-1))
-
-    scores = torch.cat(scores, dim=0)
+    with torch.no_grad():
+        for batch in val_dataloader:
+            actor_output = actor(batch)
+            cost = actor_output['total_time'].view(-1)
+            scores.append(cost)
 
-    return scores
+    return torch.cat(scores, dim=0)