Thanks ptrblck. In my machine, it’s always 3 batches, but in another machine that has the same hardware, it’s 33 batches. Today, I change the model.py and then turns to 40 batches in my machine. I haven’t store any tensors
def forward(self, input_id: torch.LongTensor,
                token_type_id: torch.LongTensor,
                attention_mask: torch.LongTensor,
                length: Tuple[int]) -> torch.FloatTensor:
        if self.use_bert_predictions:
            out, _ = self.bert(input_ids=input_id,
                               token_type_ids=token_type_id,
                               attention_mask=attention_mask)
        else:
            _, out = self.bert(input_ids=input_id,
                               token_type_ids=token_type_id,
                               attention_mask=attention_mask)
        del _
        out = out.split_with_sizes(length)
        seq_lengths = torch.LongTensor([x for x in map(len, out)])
        out = torch.nn.utils.rnn.pad_sequence(
                            out, batch_first=True)
        out = torch.nn.utils.rnn.pack_padded_sequence(
                        out, seq_lengths.cpu().numpy(), batch_first=True, enforce_sorted=False)
        out, _ = self.encoder(out)
        del _
        del seq_lengths
        out, index = torch.nn.utils.rnn.pad_packed_sequence(out, batch_first=True)
        assert len(out) == len(index)
        temp = []
        for i in range(len(out)):
            temp.append(out[i][index[i]-1])
        out = torch.stack(temp)
        del temp
        gc.collect()
        out = self.fc1(out)
        out = F.relu(out)
        out = self.fc2(out)
        return out