text: from torch import nn, Tensor
  
  
  class RobustCrossEntropyLoss(nn.CrossEntropyLoss):
      """
      this is just a compatibility layer because my target tensor is float and has an extra dimension
      """
      def forward(self, input: Tensor, target: Tensor) -> Tensor:
          if len(target.shape) == len(input.shape):
              assert target.shape[1] == 1
              target = target[:, 0]
          return super().forward(input, target.long())
      
  class MyRobustCrossEntropyLoss(nn.CrossEntropyLoss):
      """
      this is just a compatibility layer because my target tensor is float and has an extra dimension
      """
      def __init__(self):
          super(MyRobustCrossEntropyLoss, self).__init__(reduction='none')
          
      def forward(self, input: Tensor, target: Tensor) -> Tensor:
          if len(target.shape) == len(input.shape):
              assert target.shape[1] == 1
              target = target[:, 0]
          return super().forward(input, target.long())