No, this:
def my_hook(grad):
  print("In the hook")
  print(grad)

output = model(inputs)

loss = loss_function(output, labels)
model.weight_transform.register_hook(my_hook)
loss.backward()  # The hook will be called in the middle of the backward pass