Update app.py

app.py

@@ -22,7 +22,7 @@ from model import Generator, GeneratorConfig
 cfg = GeneratorConfig()
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 generator = Generator(cfg).to(device)
-generator.load_state_dict(torch.load(generator_path))
+generator.load_state_dict(torch.load(generator_path, map_location=device))
 generator.eval()
 
 # Transformations