Update path

Nested/trainers/BaseTrainer.py

@@ -3,6 +3,7 @@ import torch
 import logging
 import natsort
 import glob
+from huggingface_hub import hf_hub_download, snapshot_download
 
 logger = logging.getLogger(__name__)
 
@@ -114,5 +115,11 @@ class BaseTrainer:
         logger.info("Loading checkpoint %s", checkpoint_path)
 
         device = None if torch.cuda.is_available() else torch.device('cpu')
-        checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+        # checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+        repo_path = snapshot_download(repo_id="SinaLab/Nested")
+
+        model_file = os.path.join(repo_path, "checkpoints", "checkpoint_2.pt")
+
+        checkpoint = torch.load(model_file, map_location=device, weights_only=False)
+
         self.model.load_state_dict(checkpoint["model"])

app.py

@@ -28,6 +28,8 @@ encoder = AutoModel.from_pretrained(pretrained_path).eval()
 #     filename="checkpoints/checkpoint_2.pt"
 # )
 
+
+
 checkpoint_path = snapshot_download(repo_id="SinaLab/Nested", allow_patterns="checkpoints/")
 print("checkpoint_path : ", checkpoint_path)