squeeze output map

alexriedel1 · Jun 21, 2023 · 7188bf7 · 7188bf7
1 parent b84d753
commit 7188bf7
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/anomalib/models/components/base/anomaly_module.py b/src/anomalib/models/components/base/anomaly_module.py
@@ -179,7 +179,7 @@ def _collect_outputs(
             image_metric.update(output["pred_scores"], output["label"].int())
             if "mask" in output.keys() and "anomaly_maps" in output.keys():
                 pixel_metric.cpu()
-                pixel_metric.update(output["anomaly_maps"], output["mask"].int())
+                pixel_metric.update(torch.squeeze(output["anomaly_maps"]), output["mask"].int())
 
     @staticmethod
     def _post_process(outputs: STEP_OUTPUT) -> None: