I wrote code to inherit from PretrainedModel which adds an additional input named “formatting_features”. I trained my model successfully and saved the result using Trainer. When I load the model using AutoModel.from_pretrained(), I don’t get any errors on load, but when I attempt to use it to predict on new data, I get the following error:
TypeError: forward() got an unexpected keyword argument 'formatting_features'
When I run print(model), I only see the base LongformerModel which is contained within my custom model, but not any of the layers I added after it.
What am I doing wrong?