divide penalty by n_batches

FluxML · Sep 22, 2022 · 98394e1 · 98394e1
1 parent 61c3801
commit 98394e1
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/core.jl b/src/core.jl
@@ -36,7 +36,7 @@ function train!(loss, penalty, chain, optimiser, X, y)
         parameters = Flux.params(chain)
         gs = Flux.gradient(parameters) do
             yhat = chain(X[i])
-            batch_loss = loss(yhat, y[i]) + penalty(parameters)
+            batch_loss = loss(yhat, y[i]) + penalty(parameters)/n_batches
             training_loss += batch_loss
             return batch_loss
         end
@@ -96,7 +96,7 @@ function  fit!(loss, penalty, chain, optimiser, epochs, verbosity, X, y)
 
     parameters = Flux.params(chain)
     losses = (loss(chain(X[i]), y[i]) +
-              penalty(parameters) for i in 1:n_batches)
+              penalty(parameters)/n_batches for i in 1:n_batches)
     history = [mean(losses),]
 
     for i in 1:epochs