I’m attempting to wrap my pretrained BERT into an sklearn model with very limited success. I’ve managed to load input embeddings and labels using datasets, but I hit errors trying to pass more than one X to include the input masks and segment_ids. Any thoughts? Have people had success with this approach?