diff --git a/modules/deeplearning/cloudheight_mgpus.py b/modules/deeplearning/cloudheight_mgpus.py
index fe4d3d5c18068af7803d6b84c9053fda41282f83..65c8879fbea4359dc7cd4e9b577e216491e6f737 100644
--- a/modules/deeplearning/cloudheight_mgpus.py
+++ b/modules/deeplearning/cloudheight_mgpus.py
@@ -764,7 +764,7 @@ class CloudHeightNN:
 
             for abi, temp, lbfp, sfc in self.train_dataset:
                 trn_ds = tf.data.Dataset.from_tensor_slices((abi, temp, lbfp, sfc))
-                trn_ds = trn_ds.batch(GLOBAL_BATCH_SIZE)
+                trn_ds = trn_ds.batch(BATCH_SIZE)
                 trn_dist_ds = self.strategy.experimental_distribute_dataset(trn_ds)
                 for mini_batch in trn_dist_ds:
                     if self.learningRateSchedule is not None:
@@ -782,7 +782,7 @@ class CloudHeightNN:
 
                         for abi_tst, temp_tst, lbfp_tst, sfc_tst in self.test_dataset:
                             tst_ds = tf.data.Dataset.from_tensor_slices((abi_tst, temp_tst, lbfp_tst, sfc_tst))
-                            tst_ds = tst_ds.batch(GLOBAL_BATCH_SIZE)
+                            tst_ds = tst_ds.batch(BATCH_SIZE)
                             tst_dist_ds = self.strategy.experimental_distribute_dataset(tst_ds)
                             for mini_batch_test in tst_dist_ds:
                                 self.distributed_test_step(mini_batch_test)
@@ -810,7 +810,7 @@ class CloudHeightNN:
             self.test_accuracy.reset_states()
             for abi, temp, lbfp, sfc in self.test_dataset:
                 ds = tf.data.Dataset.from_tensor_slices((abi, temp, lbfp, sfc))
-                ds = ds.batch(GLOBAL_BATCH_SIZE)
+                ds = ds.batch(BATCH_SIZE)
                 tst_dist_ds = self.strategy.experimental_distribute_dataset(ds)
                 for mini_batch_test in tst_dist_ds:
                     self.distributed_test_step(mini_batch_test)