Hi Tom,
Thanks for your reply. In my implementation, I use a combined loss, dice + cross entropy. I suppose if the dice loss is not causing the problem, I wonder if the CE loss is. According to the documentation, I can call the CE loss in the below manner on volumetric outputs and ground truths. Is this call to cross entropy also okay?
class CombinedLoss3D(_Loss):
    A combination of dice  and cross entropy loss for volumetric outputs and ground truths

    def __init__(self, weight):
        super(CombinedLoss3D, self).__init__()
        self.cross_entropy_loss = nn.CrossEntropyLoss(weight=weight)
        self.dice_loss = DiceLoss() 

    def forward(self, output, ground_truth, weight=None):
        """
        Forward pass

        :param output: torch.tensor (NxCxDxHxW) Network output (logits) not normalized. 
        :param ground_truth: torch.tensor (NxDxHxW) 
        :param weight: torch.tensor (N) 
        :return: scalar
       """
        y_2 = self.dice_loss(output, ground_truth, weight)
        y_1 = self.cross_entropy_loss.forward(output, ground_truth)

        return y_1 + y_2

Thank you!