diffusion_brain

Runtime error

Warvito commited on Sep 16, 2022

Commit

6d6f048

1 Parent(s): 1888bc8

Changing to GPU! Thanks Omar! ^^

Files changed (1) hide show

app.py CHANGED Viewed

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

 mediapy.set_ffmpeg(ffmpeg_path)
 # Loading model
+device = torch.device("cuda")
 vqvae = mlflow.pytorch.load_model(
     "./trained_models/vae/",
     map_location=device,