fix

inference.py

@@ -21,6 +21,7 @@ def inference1():
 def inference():
   model = Unet(
     image_channels=3,
+    dropout=0.1,
   )
   model = ConditionalUnet(
     unet=model,
@@ -33,7 +34,7 @@ def inference():
     noise_steps=1000,
     beta_0=1e-4,
     beta_T=0.02,
-    image_size=(120, 80),
+    image_size=(192, 128),
   )
 
   model.to(device)

unet.py

@@ -411,7 +411,9 @@ class ConditionalUnet(nn.Module):
 
         self.class_embedding = nn.Embedding(num_classes + 1, unet.starting_channels, padding_idx=0)
     
-    def forward(self, x, t, cond=torch.tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])):
+    def forward(self, x, t, cond=None):
+        cond = torch.tensor([0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0])
+        cond = cond.unsqueeze(0)
         # cond: (batch_size, n), where n is the number of classes that we are conditioning on
         t = self.unet.time_encoding(t)