Hi @ptrblck
There has been an update, I included the ignore index in loss
criterion = nn.CrossEntropyLoss(ignore_index=0)

now my model has predictions for both class 0 and 1 during and after training. Though, I am waiting to run it for more epochs to see a final result.
Thanks