Update custom_model.py

custom_model.py

@@ -13,8 +13,12 @@ class CustomModel(PreTrainedModel):
 
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs):
-        config = cls.config_class.from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
-        model = cls(config)
-        # Load the weights from the pretrained model
-        model.load_state_dict(torch.load(pretrained_model_name_or_path))
-        return model
+        try:
+            config = cls.config_class.from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+            model = cls(config)
+            # Load the weights from the pretrained model
+            model.load_state_dict(torch.load(pretrained_model_name_or_path))
+            return model
+        except FileNotFoundError:
+            print(f"Failed to load model from {pretrained_model_name_or_path}. Please check the model name and path.")
+            return None