diff --git a/innofw/core/datasets/semantic_segmentation/tiff_dataset.py b/innofw/core/datasets/semantic_segmentation/tiff_dataset.py
index 8093d3db..57b355c8 100644
--- a/innofw/core/datasets/semantic_segmentation/tiff_dataset.py
+++ b/innofw/core/datasets/semantic_segmentation/tiff_dataset.py
@@ -109,11 +109,11 @@ def __getitem__(self, index) -> dict:
                 # out = self.transform(image=image, mask=mask)
                 # image, mask = out["image"], out["mask"]
                 image, mask = self.transform(image, mask)
-
-        try:
-            image = np.moveaxis(image, 2, 0)  # todo: refactor
-        except:
-            pass
+        else:
+            try:
+                image = np.moveaxis(image, 2, 0) 
+            except:
+                pass
         try:
             image = image.astype(np.float32)  # todo: refactor
         except: