minor change in PyTorch Wrapper

fzi-forschungszentrum-informatik · Aug 23, 2024 · a6e7b00 · a6e7b00
1 parent dcaf0e6
commit a6e7b00
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/TSInterpret/Models/PyTorchModel.py b/TSInterpret/Models/PyTorchModel.py
@@ -30,7 +30,10 @@ def predict(self, item) -> List:
         else:
             item = torch.from_numpy(item)
         out = self.model(item.float())
-        y_pred = torch.nn.functional.softmax(out).detach().numpy()
+        if out.shape[-1]>1:
+            y_pred = torch.nn.functional.softmax(out).detach().numpy()
+        else: 
+            y_pred=out.detach().numpy()
         return y_pred
 
     def load_model(self, path):