It appears there was a size mismatch I overlooked between the network output and the target. (8x2x68x68 vs 8x70x70)
I think there might be a bug there as it gives an error when the tensors are on cpu, but is silent and outputs something when they are on gpu. Do you want me to open an issue ?
That might be where the nans were coming from. I’ll let it train overnight and see.
To reproduce the error:
import torch
import torch.nn as nn
from torch.autograd import Variable

loss = nn.NLLLoss2d()

target = Variable(torch.Tensor(8, 70, 70).random_(0, 1)).long()
output = Variable(torch.randn(8, 2, 68, 68))

loss(output, target)


RuntimeError: Assertion `input0 == target0 && input2 == target1 && input3 == target2' failed. size mismatch (got input: 8x2x68x68, target: 8x70x70)


loss(output.cuda(), target.cuda())


Variable containing:
1.00000e-03 *
 -2.4779