hszhao / semseg

Semantic Segmentation in Pytorch
MIT License
1.34k stars 244 forks source link

"out of memory" during validation #60

Open domYuki opened 4 years ago

domYuki commented 4 years ago

When performing validation for each Epoch, "out of memory" occurs after validation.

You should add "with torch.no_grad():" before "output = model(input)" in tool/train.py on line 359.

Before

    model.eval()
    end = time.time()
    for i, (input, target) in enumerate(val_loader):
        data_time.update(time.time() - end)
        input = input.cuda(non_blocking=True)
        target = target.cuda(non_blocking=True)
        output = model(input)

Fixed

    model.eval()
    end = time.time()
    for i, (input, target) in enumerate(val_loader):
        data_time.update(time.time() - end)
        input = input.cuda(non_blocking=True)
        target = target.cuda(non_blocking=True)
        with torch.no_grad():
               output = model(input)