Hi,
The problem is that you modify x in place while it’s value is needed to compute gradients.
You can replace your function with:
def x_normalization(x):
       x = x / x.max(0, keepdim=True)[0]
       return x