Update custom_model.py

custom_model.py

@@ -27,8 +27,8 @@ class CustomModel(PreTrainedModel):
             config = cls.config_class.from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
             # Initialize the model with the configuration
             model = cls(config)
-            # Load the model weights using the transformers library
-            model.encoder = AutoModel.from_pretrained(pretrained_model_name_or_path, *args, **kwargs)
+            # Optionally, you can load the state_dict here if needed
+            # model.load_state_dict(torch.load(os.path.join(pretrained_model_name_or_path, "pytorch_model.bin")))
             return model
         except Exception as e:
             print(f"Failed to load model from {pretrained_model_name_or_path}. Error: {e}")