From b16b755d55728fafa953b8ac7a0f9a3916912437 Mon Sep 17 00:00:00 2001
From: torzdf <36920800+torzdf@users.noreply.github.com>
Date: Thu, 23 Oct 2025 18:18:39 +0100
Subject: [PATCH] lr-finder: Log error on NaN

---
 lib/training/lr_finder.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/lib/training/lr_finder.py b/lib/training/lr_finder.py
index d7c3298..2ed81a3 100644
--- a/lib/training/lr_finder.py
+++ b/lib/training/lr_finder.py
@@ -132,7 +132,8 @@ class LearningRateFinder:
         for idx in pbar:
             model_inputs, model_targets = self._feeder.get_batch()
             loss: list[float] = self._model.model.train_on_batch(model_inputs, y=model_targets)
-            if np.isnan(loss[0]):
+            if any(np.isnan(x) for x in loss):
+                logger.warning("NaN detected! Exiting early")
                 break
             self._on_batch_end(idx, loss[0])
             self._update_description(pbar)