That’s exactly what I’m using (pytorch-pretrained-bert is just fabulous) ! I just realized that I need to use only the last hidden state of the BertModel for single sequence classification. So, I think can replace:
model(train_batch)

with
model(train_batch)[-1]

Also, I am calculating the BERT features first because it was performing badly and I read that it can boost performance. (it’s just a workaround for some autograd problem in pytorch).