Hi @ptrblck
I have a set of 3-ch RGB images from the Kaggle DSTL Satellite Imagery Feature detection dataset. In it is has labelled polygonal data for 10 classes, ranging from buildings to trees, roads, waterways and vehicles.
I generated a binary mask for each of the classes, where each pixel value corresponds to a 1 if it contains a building or a 0 if it does not. I’ve stacked the individual marks into a single 10-ch mask.
Both the 3-ch image and 10-ch mask have been converted from numpy to torch float tensors, in the Dataset class.
For the NN model, I am using a U-Net model from https://github.com/minerva-ml/steppy-toolkit/blob/master/steppy_toolkit/pytorch/architectures/unet.py 37
and configured the U-Net model to output 10 channels:
class UNet(nn.Module):
    def __init__(self, conv_kernel=3,
                 pool_kernel=3, pool_stride=2,
                 repeat_blocks=2, n_filters=8,
                 batch_norm=True, dropout=0.1,
                 in_channels=3, out_channels=10,
                 kernel_scale=3,
                 **kwargs):

I’m then passing the output_batch and labels_batch to the multi-class cross entropy loss function in my training loop, which is where I am getting the nan error.