I am using bert for next-sentence prediction on a cpu. I want to call the model twice in a row to select a sentence from a list of sentence pairs. I am using a batch size of 128. My code looks something like this:
def bert_batch_compare(self, prompt1, prompt2):
    encoding = self.tokenizer(prompt1, prompt2, return_tensors='pt', padding=True, truncation=True, add_special_tokens=True)
    target = torch.ones((1,len(prompt1)), dtype=torch.long)
    outputs = self.model(**encoding, next_sentence_label=target)
    logits = outputs.logits.detach()
    return logits
    
def call_bert(self):
    batch_pattern = []
    batch_template = []
    batch_input = []
    ## make batch_pattern, batch_input, patch_template here!
    si = self.bert_batch_compare(batch_pattern, batch_input)
    ## second call to this fn is killed because of memory limitations    
    sj = self.bert_batch_compare(batch_input, batch_template)

If I lower the batch size to something like 24 it runs, but I’d like to use a larger batch size. I am not doing any training right now. I’m using ‘bert-base-uncased’. During the second call to ‘bert_batch_compare()’ the memory usage increases to 100% and the program crashes. I have 16G to work with. Until that time the code only uses 1.8Gig. I am using linux and python 3.6, along with pytorch 1.8.