Pass in vmap_pred as a config argument

Fung-Lab · Feb 13, 2024 · 149e7bf · 149e7bf
1 parent b3895b3
commit 149e7bf
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/matdeeplearn/tasks/task.py b/matdeeplearn/tasks/task.py
@@ -75,6 +75,7 @@ def run(self):
             # if isinstance(self.trainer.data_loader, list):
             self.trainer.predict(
                 loader=self.trainer.data_loader, split="predict", results_dir=results_dir, labels=self.config["task"]["labels"],
+                vmap_pred = self.config["task"].get("vmap_pred", False)
             )
             # else:
                 # self.trainer.predict(