Merge pull request #38 from ZerenLong/fix_bug_in_augmentation

[Bug] Fix bug in augmentations.py
visinf · May 28, 2021 · 4d7f6aa · 4d7f6aa
2 parents 9103839 + 8bec0ea
commit 4d7f6aa
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/augmentations.py b/augmentations.py
@@ -271,7 +271,7 @@ def apply_random_transforms_to_params(self,
 
             theta_transform = torch.cat([b1, b2, b3, b4, b5, b6], dim=1)
             theta_try = apply_transform_to_params(theta0, theta_transform)
-            thetas = invalid.float() * theta_try + (1 - invalid).float() * thetas
+            thetas = invalid.float() * theta_try + (1 - invalid.float()) * thetas
 
             # compute new invalid ones
             invalid = self.find_invalid(width=width, height=height, thetas=thetas)
@@ -796,7 +796,7 @@ def apply_random_transforms_to_params(self,
 
             theta_transform = torch.cat([b1, b2, b3, b4, b5, b6], dim=1)
             theta_try = apply_transform_to_params(theta0, theta_transform)
-            thetas = invalid.float() * theta_try + (1 - invalid).float() * thetas
+            thetas = invalid.float() * theta_try + (1 - invalid.float()) * thetas
 
             # compute new invalid ones
             invalid = self.find_invalid(width=width, height=height, thetas=thetas)
@@ -1075,7 +1075,7 @@ def apply_random_transforms_to_params(self,
 
             theta_transform = torch.cat([b1, b2, b3, b4, b5, b6], dim=1)
             theta_try = apply_transform_to_params(theta0, theta_transform)
-            thetas = invalid.float() * theta_try + (1 - invalid).float() * thetas
+            thetas = invalid.float() * theta_try + (1 - invalid.float()) * thetas
 
             # compute new invalid ones
             invalid = self.find_invalid(width=width, height=height, thetas=thetas)