vq-vae-2代码填坑之CUDA out of memory和element 0 of tensors does not require grad and...

运行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():
            '''
            在此处添加