def max_logit_loss(self, prediction: "torch.Tensor", target: "torch.Tensor") -> "torch.Tensor" | None:
        """
        Computes the loss as the largest logit value amongst the incorrect classes.

        :param prediction: model predictions.
        :param target: target classes. NB not one hot.
        :return: scalar loss value
        """
        target_logit = prediction[:, target]
        output = prediction - target_logit

        ubs = torch.sum(torch.abs(output[1:, :]), dim=0) + output[0, :]

        loss = None
        for i in range(self.nb_classes):
            if i != target and (loss is None or ubs[i] > loss):
                loss = ubs[i]
        return loss

        Computes the loss as the largest logit value amongst the incorrect classes.

        :param prediction: model predictions.
        :param target: target classes. NB not one hot.
        :return: scalar loss value
        
art/estimators/certification/deep_z/pytorch.py