Hi, I have a training code that runs on a single GPU like this:
def train():
    model.train()
    torch.cuda.reset_max_memory_allocated()

    for batch_idx, (inputs, targets) in enumerate(train_loader):
        inputs = inputs.cuda(non_blocking=True)
        targets = targets.cuda(non_blocking=True)

        outputs = model(inputs)
        loss = criterion(outputs, targets)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        print(f'Max memory allocated: {torch.cuda.max_memory_allocated() // 1024 // 1024}')
        torch.cuda.reset_max_memory_allocated()


def main():
    model = Net().cuda()
    cudnn.benchmark = True

    train_dataset = MyDataset()
    train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=32, shuffle=True, pin_memory=True)

    optimizer = torch.optim.SGD(model.parameters(), lr=lr)
    criterion = nn.CrossEntropyLoss()

    print(f'Max memory before training: {torch.cuda.max_memory_allocated() // 1024 // 1024}')
    train(train_loader, model, optimizer, criterion)

Then the output I got in this way is:
Max memory before training: 487
Max memory allocated: 15402
Max memory allocated: 13615
Max memory allocated: 13591
Max memory allocated: 13591
Max memory allocated: 13591
Max memory allocated: 13591
...

What could be the reason of that the GPU memory usage of the first batch is larger than the following ones? This would lead to OOM if I use slightly larger model. Is there any ways to reduce the extra usage of memory?