fix device issue

inference.py

@@ -11,7 +11,8 @@ from unet import Unet, ConditionalUnet
 
 from diffusion import GaussianDiffusion, DiffusionImageAPI
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+#device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = torch.device("mps")
 
 def inference1():
   # new image from web page

unet.py

@@ -414,6 +414,7 @@ class ConditionalUnet(nn.Module):
     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 = cond.to(self.device)
         # cond: (batch_size, n), where n is the number of classes that we are conditioning on
         t = self.unet.time_encoding(t)