Merge pull request #8 from geoaigroup/aghand0ur-patch-1

Reverting to intial TorchRandomRotate as in NasaHarvest
geoaigroup · Dec 18, 2023 · d911639 · d911639
2 parents 2887fb2 + 219826e
commit d911639
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/data_processing/augmentation.py b/data_processing/augmentation.py
@@ -50,9 +50,8 @@ def get_params(degrees) -> float:
     @torch.no_grad()
     def __call__(self,img,mask=None):
 
-        # batch_size = img.shape[0]
-        batch_size=img.width
-        print(batch_size)
+        batch_size = img.shape[0]
+
         for i in range(batch_size):
 
             if random.random() > self.proba: