Switch over to using torch.testing.assert_close in tests

[epwalsh]

assert_allclose is being deprecated. See pytorch/pytorch#61844.

[dirkgr]

The suggested replacement was only added in torch 1.9.0, so we need to wait until that is the oldest version we support before we tackle this.

[epwalsh]

I don't think this is a user-facing change but let me know if I'm wrong. Because as far as I know we only use this in internal tests, and we only run those tests against the latest version of PyTorch.