Upload app.py

app.py

@@ -83,7 +83,7 @@ def get_paths(dataset_id):
 
 
 def load_pgm(dataset_id, pgm_path):
-    checkpoint = torch.load(pgm_path, map_location=DEVICE, weights_only=True)
+    checkpoint = torch.load(pgm_path, map_location=DEVICE)
     args = Hparams()
     args.update(checkpoint["hparams"])
     args.device = DEVICE
@@ -101,7 +101,7 @@ def load_pgm(dataset_id, pgm_path):
 def load_vae(dataset_id, vae_path):
     if "Chest" in dataset_id:
         vae_path, dscm_path = vae_path[0], vae_path[1]
-    checkpoint = torch.load(vae_path, map_location=DEVICE, weights_only=True)
+    checkpoint = torch.load(vae_path, map_location=DEVICE)
     args = Hparams()
     args.update(checkpoint["hparams"])
     # backwards compatibility hack
@@ -115,7 +115,7 @@ def load_vae(dataset_id, vae_path):
     vae = HVAE(args).to(args.device)
 
     if "Chest" in dataset_id:
-        dscm_ckpt = torch.load(dscm_path, map_location=DEVICE, weights_only=True)
+        dscm_ckpt = torch.load(dscm_path, map_location=DEVICE)
         vae.load_state_dict(
             {
                 k[4:]: v