diff --git a/osrt/encoder.py b/osrt/encoder.py
index 2388dd34a7007444d498167abcf0fe577d405c9b..cc64f4a4aec11c06a83da06b1d0f78f079b7910a 100644
--- a/osrt/encoder.py
+++ b/osrt/encoder.py
@@ -188,7 +188,7 @@ class FeatureMasking(nn.Module):
         B, N, H, W, C  = images.shape
         orig_size = (H, W)
         images = images.reshape(B*N, H, W, C) # [B x N, C, H, W]         
-        im_size = self.resize(images[0]).shape[-3:-1]
+        im_size = self.resize.apply_image(images[0]).shape[-3:-1]
 
         ### Pre-process images for the image encoder (Resize and Pad)
         images = torch.stack([self.preprocess(x) for x in images])