From c96c13b65576d954e8413140c0bd40637e0450f4 Mon Sep 17 00:00:00 2001
From: Victor Joos <victor.joos@uclouvain.be>
Date: Tue, 28 Jun 2022 16:02:29 +0200
Subject: [PATCH] fix visu

---
 hydra_flash/train.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/hydra_flash/train.py b/hydra_flash/train.py
index 6214831..90fc124 100644
--- a/hydra_flash/train.py
+++ b/hydra_flash/train.py
@@ -30,7 +30,9 @@ def main(cfg):
     predict_datamodule = instantiate(cfg.predict_datamodule)
     predictions = trainer.predict(model, datamodule=predict_datamodule, output="labels")
 
-    show_predictions(predict_datamodule.predict_dataset, predictions[0])
+    show_predictions(
+        predict_datamodule.predict_dataset, predictions[0], cfg.show, cfg.save
+    )
 
 
 if __name__ == "__main__":
-- 
GitLab