Update archs/model.py

archs/model.py

@@ -33,6 +33,8 @@ class AttentionBlock(nn.Module):
 class UNet(nn.Module):
     def __init__(self):
         super(UNet, self).__init__()
+        
+        self.padder_size = 32
 
         self.encoder = nn.Sequential(
             nn.Conv2d(3, 32, kernel_size=3, padding=1),
@@ -68,6 +70,10 @@ class UNet(nn.Module):
         )
 
     def forward(self, x):
+
+        _, _, H, W = x.shape
+        x = self.check_image_size(x)
+        
         skip_connections = []
 
         for layer in self.encoder:
@@ -93,5 +99,12 @@ class UNet(nn.Module):
             else:
                 x = layer(x)
 
-        return x
+        return x[:, :, :H, :W]
 
+        
+    def check_image_size(self, x):
+        _, _, h, w = x.size()
+        mod_pad_h = (self.padder_size - h % self.padder_size) % self.padder_size
+        mod_pad_w = (self.padder_size - w % self.padder_size) % self.padder_size
+        x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), value = 0)
+        return x