Hi,
I get the cuda runtime error (2) during training. However, it’s strange as the 1st iteration runs ok, but the 2nd iteration returns the runtime error. My model has 32 million parameters, I have 2 nVidia 1080Ti which means 22GB of GPU memory. Feeding images of 256x256x3 with a batch size of 40.
Here’s the code:
for epoch in range(1, epochs):
    train_loss = 0.0
    for i, batch in enumerate(loader):
        ims, masks = batch['image'].cuda(), batch['mask'].cuda().float()

        ims = Variable(ims)
        masks = Variable(masks)

        optimizer.zero_grad()

        outputs = net(ims)

        loss = criterion(outputs, masks)
        loss.backward()
        optimizer.step()

        train_loss += loss.data[0]

I tried adding torch.cuda.empty_cache() and that didn’t help.
Any advice on how to fix this? Why does it run the 1st iteration and run out of GPU memory for the 2nd iteration?