debug

revq/models/revq.py

@@ -36,11 +36,11 @@ class Viewer:
 
 class ReVQ(PyTorchModelHubMixin, nn.Module):
     @classmethod
-    def _from_pretrained(cls, repo_id: str, **kwargs):
-        print(f"Loading ReVQ model from {repo_id}...")
-        config_path = hf_hub_download(repo_id=repo_id, filename="512T_NC=16384.yaml")
-        ckpt_path = hf_hub_download(repo_id=repo_id, filename="ckpt.pth")
-        
+    def _from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
+        print(f"Loading ReVQ model from {pretrained_model_name_or_path}...")
+        config_path = hf_hub_download(repo_id=pretrained_model_name_or_path, filename="512T_NC=16384.yaml")
+        ckpt_path = hf_hub_download(repo_id=pretrained_model_name_or_path, filename="ckpt.pth")
+
         full_cfg = OmegaConf.load(config_path)
         model_cfg = full_cfg.get("model", {})
         decoder_config = model_cfg.get("decoder", {})