Upload app.py

app.py

@@ -36,7 +36,7 @@ def _denormalize(input: torch.Tensor) -> torch.Tensor:
 
 
 def generate_images(generator, number_imgs):
-    noise = torch.zeros(number_imgs, 256, 1, 1, device='cuda').normal_(0.0, 1.0)
+    noise = torch.zeros(number_imgs, 256, 1, 1, device=device).normal_(0.0, 1.0)
     with torch.no_grad():
         gan_images, _ = generator(noise)
 
@@ -127,7 +127,7 @@ def main():
                     <small><i>Please stand by.</i></small>
                 """,
                 unsafe_allow_html=True,)
-                
+
         if generate_button:
             generator = load_generator(model_name[img_type])
             gan_images = generate_images(generator, number_imgs)