From e73cc0bdd61c5ba6bfca4647bfbb0f4eedbd3540 Mon Sep 17 00:00:00 2001
From: SCiarella <simoneciarella@gmail.com>
Date: Sun, 3 Nov 2024 14:56:00 +0100
Subject: [PATCH] Add warning for missing val_loss

---
 src/speckcn2/loss.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/src/speckcn2/loss.py b/src/speckcn2/loss.py
index 082a938..a3bef84 100644
--- a/src/speckcn2/loss.py
+++ b/src/speckcn2/loss.py
@@ -57,7 +57,11 @@ def __init__(self,
                  validation: bool = False):
         super(ComposableLoss, self).__init__()
         if validation:
-            config['loss'] = config['val_loss']
+            if 'val_loss' in config:
+                config['loss'] = config['val_loss']
+            else:
+                print('[!] Warning: Validation loss not found in config.yaml,',
+                      'keeping track of training loss instead')
         self.device = device
         self.loss_functions: dict[str, Callable] = {
             'MSE': torch.nn.MSELoss(),