Thank you very much the pointer, I was able to redo do the function and now it runs (almost) as fast as the original softmax! For reference here is the function (used for a 3D tensor and calculated over the rows):
    def forward(self, input):
        m = 0
        expm = math.exp(-m) #calculate m
        input = torch.exp(input)
        
        #keep track of the original dimensions
        s0 = input.size(0)
        s1 = input.size(1)
        s2 = input.size(2)
        
        #input = input.transpose(1,2).contiguous().view(-1,s1) #use if you want to calculate over columns instead of rows
        input = input.view(-1,s2)
        
        #calculate softmargin softmax
        expsum = torch.sum(input, dim=1).view(-1,1) # sum and reshape for softmax
        input = (input*expm)/(expsum - input + (input*expm))
        
        #normalize
        normsum = torch.sum(input, dim=1).view(-1,1) #sum and reshape for normalization
        input = input/normsum
        
        input = input.view(s0,s1,s2)
        #input = input.view(s0,s2,s1).transpose(1,2).contiguous() #  #use if you want to calculate over columns instead of rows
        
        return input