As I see their code, something like that might do the trick.
def sparsemax(input, dim):
    """
    Args:
        input (torch.Tensor): Input tensor. First dimension should be the batch size
    Returns:
        torch.Tensor: [batch_size x number_of_logits] Output tensor
    """
    # Sparsemax currently only handles 2-dim tensors,
    # so we reshape to a convenient shape and reshape back after sparsemax
    input = input.transpose(0, dim)
    original_size = input.size()
    input = input.reshape(input.size(0), -1)
    input = input.transpose(0, 1)
    dim = 1
    ...
    ...