Update app.py

app.py

@@ -119,9 +119,9 @@ class UNet(nn.Module):
         x = self.up3(x, x1, t)
         output = self.outc(x)
         return output
-device = 'cpu'
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
 model = UNet(device = device).to(device)
-model.load_state_dict(torch.load('Model_Saved_States/diffusion_64.pth'))
+model.load_state_dict(torch.load('Model_Saved_States/diffusion_64.pth', map_location=torch.device(device)))
 img_size = 64
 class Diffusion():
   def __init__(self, time_steps = 500, beta_start = 0.0001, beta_stop = 0.02, image_size = 64, device = device):