Update model.py

model.py

@@ -18,20 +18,23 @@ class ModelColorization(nn.Module, PyTorchModelHubMixin):
             nn.ReLU(),
             nn.BatchNorm2d(64),
             nn.Flatten(),
-            nn.Linear(64 * 16 * 16, 4000),
+            nn.Linear(64*16*16, 4000),
         )
         self.decoder = nn.Sequential(
             nn.Linear(4000, 64 * 16 * 16),
             nn.ReLU(),
+
             nn.Unflatten(1, (64, 16, 16)),
             nn.ConvTranspose2d(64, 128, kernel_size=2, stride=2),
             nn.ReLU(),
             nn.BatchNorm2d(128),
+
             nn.ConvTranspose2d(128, 256, kernel_size=2, stride=2),
             nn.ReLU(),
             nn.BatchNorm2d(256),
+
             nn.ConvTranspose2d(256, 3, kernel_size=2, stride=2),
-            nn.Sigmoid(),
+            nn.Sigmoid()
         )
 
     def forward(self, x):