Update UNet.py

UNet.py

@@ -14,7 +14,6 @@ class Block(Module):
     def __init__(self, inChannels, outChannels):
         super().__init__()
         # store the convolution and RELU layers
-        self.height, self.width = 256, 256
         self.conv1 = Conv2d(inChannels, outChannels, 3)
         self.relu = ReLU()
         self.conv2 = Conv2d(outChannels, outChannels, 3)
@@ -73,7 +72,7 @@ class Decoder(Module):
     
 class UNet(Module):
     def __init__(self, encChannels=(3, 64, 128, 256, 512, 1024), decChannels=(1024, 512, 256, 128, 64),
-                 nbClasses=1, retainDim=True, outSize=(height=256, width=256)):
+                 nbClasses=1, retainDim=True, outSize=(256, 256)):
         super().__init__()
         # initialize the encoder and decoder
         self.encoder = Encoder(encChannels)