Update archs/model.py

archs/model.py

@@ -1,6 +1,5 @@
 import torch
 import torch.nn as nn
-import torch.nn.functional as F
 
 class AttentionBlock(nn.Module):
     def __init__(self, in_channels, out_channels, kernel_size=3, padding=1):
@@ -30,11 +29,10 @@ class AttentionBlock(nn.Module):
         return x
 
 
+
 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),
@@ -70,10 +68,6 @@ 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:
@@ -99,12 +93,5 @@ class UNet(nn.Module):
             else:
                 x = layer(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
+