Ok, I finally fixed it, and made a clean method to add new features to any Bert model. Thanks for your help guys.