From 219826e25b7685decb866503c856f2169c0572a6 Mon Sep 17 00:00:00 2001
From: "Ali J. Ghandour" <ghandour.aub@gmail.com>
Date: Mon, 18 Dec 2023 18:36:55 +0200
Subject: [PATCH] Update augmentation.py

---
 data_processing/augmentation.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/data_processing/augmentation.py b/data_processing/augmentation.py
index e53c959..558f0ee 100644
--- 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: