Fixed identity selection mask visualisation bug

ArnavVarma · Sep 2, 2019 · 3e263c0 · 3e263c0
1 parent 1cc8b81
commit 3e263c0
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/trainer.py b/trainer.py
@@ -475,7 +475,8 @@ def compute_losses(self, inputs, outputs):
                 to_optimise, idxs = torch.min(combined, dim=1)
 
             if not self.opt.disable_automasking:
-                outputs["identity_selection/{}".format(scale)] = (idxs > 1).float()
+                outputs["identity_selection/{}".format(scale)] = (
+                    idxs > identity_reprojection_loss.shape[1] - 1).float()
 
             loss += to_optimise.mean()