运行vq-vae-2-pytorch项目,第三阶段:
python train_pixelsnail.py lmdbdata
报错RuntimeError: CUDA out of memory. Tried to allocate 1.50 GiB…
在out前加入with torch.no_grad():
触发另外一个错误:
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
源代码段为:
def train(args, epoch, loader, model, optimizer, scheduler, device):
loader = tqdm(loader)
criterion = nn.CrossEntropyLoss()
for i, (top, bottom, label) in enumerate(loader):
model.zero_grad()
top = top.to(device)
if args.hier == 'top':
target = top
with torch.no_grad():
'''
在此处添加