[Python] GPU Memory leak in pytorch

I have a weird problem in pytorch.

I have two loss, called loss_1 and loss_2. For convenience, I calculate both, but just backward loss_1. Then, it OOM. However, if I backward both loss, Loss_1 and Loss_2, it will be normal.

for epoch_id in trange(last_epoch + 1, args.n_epochs, desc="Epoch"):
model.train()
for qids, labels, *input_data in tqdm(train_dataloader, desc="Batch"):
optimizer.zero_grad()
bs = labels.size(0)
for a in range(0, bs, args.mini_batch_size):
b = min(a + args.mini_batch_size, bs)
if args.prior_loss:
logits, loss_2, _ = model(*[x[a:b] for x in input_data])
else:
logits, _ = model(*[x[a:b] for x in input_data])
# logits: [bs, nc]
loss = loss / (b - a)
if args.prior_loss and epoch_id>args.prior_loss_epoch:
loss+=loss_2
elif args.prior_loss:
del loss_2
loss.backward()
optimizer.step()
scheduler.step()

Continue reading...