Update utils/common_viz.py

utils/common_viz.py

@@ -106,23 +106,23 @@ def get_batch(
 
 def init(
     config_name: str,
-) -> Tuple[Diffuser, clip.model.CLIP, MultimodalDataset, torch.device]:
+) -> Tuple[Diffuser, clip.model.CLIP, MultimodalDataset]:
     with initialize(version_base="1.3", config_path="../configs"):
         config = compose(config_name=config_name)
 
     OmegaConf.register_new_resolver("eval", eval)
 
     # Initialize model
-    device = torch.device(config.compnode.device)
+    # device = torch.device(config.compnode.device)
     diffuser = instantiate(config.diffuser)
-    state_dict = torch.load(config.checkpoint_path, map_location=device)["state_dict"]
+    state_dict = torch.load(config.checkpoint_path, map_location="cpu")["state_dict"]
     state_dict["ema.initted"] = diffuser.ema.initted
     state_dict["ema.step"] = diffuser.ema.step
     diffuser.load_state_dict(state_dict, strict=False)
-    diffuser.to(device).eval()
+    diffuser.to("cpu").eval()
 
     # Initialize CLIP model
-    clip_model = load_clip_model("ViT-B/32", device)
+    clip_model = load_clip_model("ViT-B/32", "cpu")
 
     # Initialize dataset
     config.dataset.char.load_vertices = True
@@ -133,4 +133,4 @@ def init(
     diffuser.get_matrix = dataset.get_matrix
     diffuser.v_get_matrix = dataset.get_matrix
 
-    return diffuser, clip_model, dataset, device
+    return diffuser, clip_model, dataset, config.compnode.device