@staticmethod
def forward(ctx, *inputs):
    """ We manually put max_norm to the end of tensor inputs in order to please Function.apply """
    ctx._max_norm = inputs[-1].item()
    return inputs[:-1]
