FIX: Distribution

interTwin-eu · Aug 8, 2023 · 0db2df2 · 0db2df2
1 parent c9e5c73
commit 0db2df2
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 475 deletions.
diff --git a/ai/src/itwinai/models/tensorflow/cyclegan.py b/ai/src/itwinai/models/tensorflow/cyclegan.py
@@ -403,6 +403,9 @@ def train_step(self, batch_data):
             "D_Y_loss": disc_Y_loss,
         }
 
+    def call(self, inputs):
+        return self.gen_G(inputs, training=False)
+
     def get_config(self):
         config = super().get_config().copy()
         config.update({

diff --git a/use-cases/zebra2horse/pipeline.yaml b/use-cases/zebra2horse/pipeline.yaml
@@ -7,7 +7,7 @@ trainer:
   class_path: trainer.Zebra2HorseTrainer
   init_args:
     epochs: 10
-    batch_size: 1
+    batch_size: 4
     model:
       class_path: itwinai.models.tensorflow.cyclegan.CycleGAN
       init_args: