I am trying to train a BERT model, But one pass consumes about 11Gb. I am using the transformers library. Is the library not optimized enough or what?
Here’s the code :
from transformers import BertConfig, BertForSequenceClassification
import torch

config = BertConfig()
model =  BertForSequenceClassification(config)

model.cuda(5)
model.forward(input_ids=torch.ones([15,512],dtype=torch.int64).cuda(5),
              attention_mask=torch.ones([15,512],dtype=torch.int64).cuda(5),
              labels=torch.ones([15],dtype=torch.int64).cuda(5))

I just mentioned the relevant code here, but I can provide other snippets if needed