How can I get the sensitivity of the loss function with respect to the NN parameters?

[chan-mander]

I'm relatively new to neural networks and Brax, and I'm running into a specific issue. The example I am trying to run is below:

    for i in range(timeSteps):
        action = model.forward(observation)
        observation, reward, done, info = env.step(action) # Brax env wrapped in VectorGymWrapper and then wrapped in TorchWrapper
        observation = observation.requires_grad_()

        with torch.no_grad():
            loss += torch.pow(torch.norm(observation, p=2), 2)

    loss = loss/timeSteps
    print(loss)
    loss.backward()
    for name, param in model.named_parameters():
        print(name, param.grad)

When I run my example, the param.grad returns None. How can I get the sensitivity of the loss function with respect to the parameters?
I assume that this should be possible, because of the example that was posted here is able to do it:

https://colab.research.google.com/github/google/brax/blob/main/notebooks/training_torch.ipynb