Update modeling.py

modeling.py

@@ -114,10 +114,14 @@ class MARModel(PreTrainedModel):
     def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
         config = MARConfig.from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
         model = cls(config)
-        state_dict = torch.load('checkpoint-last.safetensors')
+        safetensors_path = os.path.join(pretrained_model_name_or_path, "checkpoint-last.safetensors")
+        if not os.path.exists(safetensors_path):
+            raise FileNotFoundError(f"safetensors file not found at {safetensors_path}")
+        state_dict = torch.load(safetensors_path, map_location='cpu')
         model.model.load_state_dict(state_dict)
         return model
 
+
     def save_pretrained(self, save_directory):
       # we will save to safetensors
       os.makedirs(save_directory, exist_ok=True)