CrossEntropyLoss on CPU is leaking #238
Description
running the imagenet example will slow leak memory
- gc.collect() does not help
- dataparallel is not the issue, even single-GPU job has the leak
This is the snippet that will leak:
for epoch in range(args.nEpochs):
for i, data in enumerate(train_loader, 0):
input, label = data
input = Variable(input)
label = Variable(label)
def closure():
output = model(input)
loss = criterion(output, label)
optimizer.zero_grad()
loss.backward()
optimizer.step(closure)
print(i)This wont leak:
for epoch in range(args.nEpochs):
for i, data in enumerate(train_loader, 0):
input, label = data
input = Variable(input.cuda())
label = Variable(label.cuda())
def closure():
output = model(input)
loss = criterion(output, label)
optimizer.zero_grad()
loss.backward()
optimizer.step(closure)
print(i)