Fix device error

app.py

@@ -19,19 +19,13 @@ ffmpeg_path = shutil.which("ffmpeg")
 mediapy.set_ffmpeg(ffmpeg_path)
 
 # Loading model
-device = torch.device("cuda")
-vqvae = mlflow.pytorch.load_model(
-    "./trained_models/vae/",
-    map_location=device,
-)
+vqvae = mlflow.pytorch.load_model("./trained_models/vae/")
 vqvae.eval()
 
-diffusion = mlflow.pytorch.load_model(
-    "./trained_models/ddpm/",
-    map_location=device,
-)
+diffusion = mlflow.pytorch.load_model("./trained_models/ddpm/")
 diffusion.eval()
 
+device = torch.device("cuda")
 diffusion = diffusion.to(device)
 vqvae = vqvae.to(device)