I am trying to train a model with the Kitti road segmentation dataset: http://www.cvlibs.net/datasets/kitti/eval_road.php 2.
I am getting Nan from the CrossEntropyLoss module. Notice that it is returning Nan already in the first mini-batch. I already checked my input tensor for Nans and Infs.
The tensor shapes I am giving to the loss func are: (b_size, n_class, h, w) and (b_size, h, w).
When I try to reshape the tensor in the following way:
loss = criterion(prediction.permute(0,2,3,1).contiguous().view(-1, n_class), target.view(-1))

the Nans disappear.
Someone knows what can be happening?
I am afraid that maybe something is still wrong, because I am still having troubles too make the network to converge.