Update app.py

app.py

@@ -9,10 +9,10 @@ device = "cpu"
 @torch.inference_mode()
 def inference_gan():
     generator = torch.jit.load("mnist-G-torchscript.pt").to(device)
-    x = torch.randn(30, 256, device='cuda')
+    x = torch.randn(30, 256, device=device)
     y = generator(x)
     y = y.view(-1, 1, 28, 28)  # reshape y to have 1 channel
-    grid = make_grid(y.cpu().detach(), nrow=8)
+    grid = make_grid(y.detach(), nrow=8)
     img = T.functional.to_pil_image(grid)
     return img
 
@@ -22,10 +22,10 @@ def inference_dcgan():
     def denorm(img_tensors):
         stats = (0.5, 0.5, 0.5), (0.5, 0.5, 0.5)
         return img_tensors * stats[1][0] + stats[0][0]
-    x = torch.randn(64, 128, 1, 1, device='cuda')
+    x = torch.randn(64, 128, 1, 1, device=device)
     y = generator(x)
     y = y.view(-1, 3, 64, 64)  # reshape y to have 3 channels
-    grid = make_grid(denorm(y.cpu().detach()), nrow=8)
+    grid = make_grid(denorm(y.detach()), nrow=8)
     img = T.functional.to_pil_image(grid)
     return img
 def inference_both():