Fix model shortcut bug (#100)

Transforms occurring before shortcut test could give false negative that shortcut should be taken
HakaiInstitute · Feb 9, 2024 · 044a9c3 · 044a9c3
1 parent f9f45ff
commit 044a9c3
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/kelp_o_matic/managers.py b/kelp_o_matic/managers.py
@@ -74,12 +74,12 @@ def __call__(self):
                 # Reorder bands
                 crop = crop[:, :, [b - 1 for b in self.band_order]]
 
-                if self.model.transform:
-                    crop = self.model.transform(crop / self._max_value)
-
-                if torch.all(crop == 0):
+                if (crop == 0).all():
                     logits = self.model.shortcut(self.reader.crop_size)
                 else:
+                    if self.model.transform:
+                        crop = self.model.transform(crop / self._max_value)
+
                     # Zero pad to correct shape
                     _, h, w = crop.shape
                     crop = torch.nn.functional.pad(