Fixed DSR (openvinotoolkit#1486)

* fixed DSR squeeze bug * added comment
ashwinvaidya17 · Jan 23, 2024 · b8bd3f9 · b8bd3f9
1 parent ee1f659
commit b8bd3f9
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/src/anomalib/models/dsr/torch_model.py b/src/anomalib/models/dsr/torch_model.py
@@ -154,6 +154,10 @@ def forward(self, batch: Tensor, anomaly_map_to_generate: Tensor | None = None)
                 ).detach()
                 image_score = torch.amax(out_mask_averaged, dim=(2, 3)).squeeze()
 
+                # prevent crash when image_score is a single value (batch size of 1)
+                if image_score.size() == torch.Size([]):
+                    image_score = image_score.unsqueeze(0)
+
                 out_mask_cv = out_mask_sm_up[:, 1, :, :]
 
                 outputs = {"anomaly_map": out_mask_cv, "pred_score": image_score}