fix: return prediction proba

instadeepai · Feb 8, 2025 · d67cc73 · d67cc73
1 parent 2f9f76e
commit d67cc73
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/instageo/model/infer_utils.py b/instageo/model/infer_utils.py
@@ -150,9 +150,8 @@ def chip_inference(
                 prediction_batch = model(data)
                 prediction_cls = (
                     torch.nn.functional.softmax(prediction_batch, dim=1)
-                    .argmax(dim=1)
                     .cpu()
-                    .numpy()
+                    .numpy()[:, 1, :, :]
                 )
 
                 profiles = []