Compatibility with Pytorch 2.0; failing test `test_gradient_value`

I'm seeing one test failure in https://github.com/conda-forge/fairscale-feedstock/pull/28 when built against pytorch 2.0

```
=================================== FAILURES ===================================
_____________________________ test_gradient_value ______________________________

def test_gradient_value():
"""Test that we don't mutate the gradients during backward"""
model = Linear(2, 2, bias=False)
optim = AdaScale(SGD(model.parameters(), lr=0.1), num_gradients_to_accumulate=2)

# fwd 1
out = model(Tensor([0.0, 1.0]))
out.sum().backward()
assert np.allclose(model.weight.grad.numpy(), [[0.0, 1.0], [0.0, 1.0]]), model.weight.grad

# fwd 2, grad is accumulated
out = model(Tensor([0.0, 1.0]))
out.sum().backward()
assert np.allclose(model.weight.grad.numpy(), [[0.0, 2.0], [0.0, 2.0]]), model.weight.grad

# assert gain and grad value before/after step/zero_grad
assert np.allclose(optim.gain(), 1.0000002499999376), optim.gain()
optim.step()
assert np.allclose(model.weight.grad.numpy(), [[0.0, 2.0], [0.0, 2.0]]), model.weight.grad
optim.zero_grad()
> assert np.allclose(model.weight.grad.numpy(), [[0.0, 0.0], [0.0, 0.0]]), model.weight.grad
E AttributeError: 'NoneType' object has no attribute 'numpy'
```

Looks like `model.weight.grad` becomes None somehow.

Start by running the failing test_gradient_value test from the linked fairscale-feedstock pull request with PyTorch 2.0, then inspect the AdaScale optimizer behavior around step() and zero_grad(). Done means the test passes and model.weight.grad remains a zero-valued tensor after zero_grad(), rather than becoming None.

Written by the indexing model from the issue text.