FIX: Distributed Dataset

interTwin-eu · Aug 9, 2023 · 1c19f66 · 1c19f66
1 parent 00d4c56
commit 1c19f66
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/ai/src/itwinai/backend/tensorflow/trainer.py b/ai/src/itwinai/backend/tensorflow/trainer.py
@@ -53,8 +53,8 @@ def train(self, data):
         train, test = data
 
         # Set batch size to the dataset
-        train = train.batch(self.batch_size * self.num_devices, drop_remainder=True)
-        test = test.batch(self.batch_size * self.num_devices, drop_remainder=True)
+        train = train.batch(self.batch_size, drop_remainder=True)
+        test = test.batch(self.batch_size, drop_remainder=True)
 
         # Number of samples
         n_train = train.cardinality().numpy()