beacon-biosignals / lighthouseflux.jl Goto Github PK

diff --git a/src/LighthouseFlux.jl b/src/LighthouseFlux.jl
index 55ccc2b..d5af80b 100644
--- a/src/LighthouseFlux.jl
+++ b/src/LighthouseFlux.jl
@@ -85,24 +85,27 @@ end

 function Lighthouse.train!(classifier::FluxClassifier, batches, logger)
     Flux.trainmode!(classifier.model)
-    weights = Zygote.Params(classifier.params)
-    for batch in batches
-        train_loss, back = log_resource_info!(logger, "train/forward_pass";
-                                              suffix="_per_batch") do
-            f = () -> loss(classifier.model, batch...)
-            return Zygote.pullback(f, weights)
-        end
-        log_value!(logger, "train/loss_per_batch", train_loss)
-        gradients = log_resource_info!(logger, "train/reverse_pass";
-                                       suffix="_per_batch") do
-            return back(Zygote.sensitivity(train_loss))
-        end
-        log_resource_info!(logger, "train/update"; suffix="_per_batch") do
-            Flux.Optimise.update!(classifier.optimiser, weights, gradients)
-            return nothing
+    try
+        weights = Zygote.Params(classifier.params)
+        for batch in batches
+            train_loss, back = log_resource_info!(logger, "train/forward_pass";
+                                                  suffix="_per_batch") do
+                f = () -> loss(classifier.model, batch...)
+                return Zygote.pullback(f, weights)
+            end
+            log_value!(logger, "train/loss_per_batch", train_loss)
+            gradients = log_resource_info!(logger, "train/reverse_pass";
+                                           suffix="_per_batch") do
+                return back(Zygote.sensitivity(train_loss))
+            end
+            log_resource_info!(logger, "train/update"; suffix="_per_batch") do
+                Flux.Optimise.update!(classifier.optimiser, weights, gradients)
+                return nothing
+            end
         end
+    finally
+        Flux.testmode!(classifier.model)
     end
-    Flux.testmode!(classifier.model)
     return nothing
 end