Update model.py

model.py

@@ -18,10 +18,10 @@ class ModelColorization(nn.Module, PyTorchModelHubMixin):
             nn.ReLU(),
             nn.BatchNorm2d(64),
             nn.Flatten(),
-            nn.Linear(64 * 16 * 16, 1024),
+            nn.Linear(64 * 16 * 16, 3000),
         )
         self.decoder = nn.Sequential(
-            nn.Linear(1024, 64 * 16 * 16),
+            nn.Linear(3000, 64 * 16 * 16),
             nn.ReLU(),
             nn.Unflatten(1, (64, 16, 16)),
             nn.ConvTranspose2d(64, 128, kernel_size=2, stride=2),