Try to fix device error

app.py

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