diff --git a/modules/deeplearning/cloudheight_mgpus.py b/modules/deeplearning/cloudheight_mgpus.py
index 65c8879fbea4359dc7cd4e9b577e216491e6f737..98210dd4bc4ebf66b1817b2393ec2dbdbe882363 100644
--- a/modules/deeplearning/cloudheight_mgpus.py
+++ b/modules/deeplearning/cloudheight_mgpus.py
@@ -669,7 +669,7 @@ class CloudHeightNN:
         with tf.GradientTape() as tape:
             pred = self.model(inputs, training=True)
             loss = self.loss(labels, pred)
-            loss = tf.nn.compute_average_loss(loss, global_batch_size=GLOBAL_BATCH_SIZE)
+            loss = tf.nn.compute_average_loss(loss, global_batch_size=BATCH_SIZE)
             total_loss = loss
             if len(self.model.losses) > 0:
                 reg_loss = tf.math.add_n(self.model.losses)