bug fix matrix dimensions increasing

Aske-Rosted · Jan 22, 2024 · 8718ff8 · 8718ff8
1 parent 93038c6
commit 8718ff8
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/src/graphnet/deployment/i3modules/graphnet_module.py b/src/graphnet/deployment/i3modules/graphnet_module.py
@@ -238,7 +238,9 @@ def _inference(self, data: Data) -> np.ndarray:
             len(task_predictions) == 1
         ), f"""This method assumes a single task. \n
                Got {len(task_predictions)} tasks."""
-        return self.model(data)[0].detach().numpy()
+        return (
+            task_predictions[0].detach().numpy()
+        )  # self.model(data)[0].detach().numpy()
 
 
 class I3PulseCleanerModule(I3InferenceModule):