Thanks for your reply!
it works fine:
def myLoss(output, y):
    loss = (output.view(512, -1) - y.view(512, -1)).abs_()
    mask = loss > threshold
    loss[mask] = loss[mask].add(1).log().mul(50)
    
    return loss.mean()