Warvito commited on
Commit
d90e254
1 Parent(s): 674eb4e
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -19,16 +19,16 @@ ffmpeg_path = shutil.which("ffmpeg")
19
  mediapy.set_ffmpeg(ffmpeg_path)
20
 
21
  # Loading model
22
- device = torch.device("cuda")
23
  vqvae = mlflow.pytorch.load_model(
24
  "./trained_models/vae/",
25
- # map_location=device,
26
  )
27
  vqvae.eval()
28
 
29
  diffusion = mlflow.pytorch.load_model(
30
  "./trained_models/ddpm/",
31
- # map_location=device,
32
  )
33
  diffusion.eval()
34
 
 
19
  mediapy.set_ffmpeg(ffmpeg_path)
20
 
21
  # Loading model
22
+ device = torch.device("cpu")
23
  vqvae = mlflow.pytorch.load_model(
24
  "./trained_models/vae/",
25
+ map_location=device,
26
  )
27
  vqvae.eval()
28
 
29
  diffusion = mlflow.pytorch.load_model(
30
  "./trained_models/ddpm/",
31
+ map_location=device,
32
  )
33
  diffusion.eval()
34