replace assert torch.allclose with torch.testing.assert_allclose

[pmeier]

Saw this while going through the test suite. We shouldn't use the assert torch.allclose idiom for a couple of reasons:

1. It broadcasts the inputs:

>>> a = torch.zeros(3, 2)
>>> b = torch.zeros(4, 3, 2)
>>> assert torch.allclose(a, b)

2. It only gives a boolean output without any diagnostic information:

>>> a[0, 0] = 1
>>> assert torch.allclose(a, b)
AssertionError

Both of these things are undesirable inside the test suite. Instead we should be using torch.testing.assert_close that has neither of the issues mentioned above:

>>> a = torch.zeros(3, 2)
>>> b = torch.zeros(4, 3, 2)
>>> torch.testing.assert_close(a, b)
AssertionError: The values for attribute 'shape' do not match: torch.Size([3, 2]) != torch.Size([4, 3, 2]).
>>> b = a.clone()
>>> a[0, 0] = 1
>>> torch.testing.assert_close(a, b)
AssertionError: Tensor-likes are not close!
Mismatched elements: 1 / 6 (16.7%)
Greatest absolute difference: 1.0 at index (0, 0) (up to 1e-05 allowed)
Greatest relative difference: inf at index (0, 0) (up to 1.3e-06 allowed)

[pmeier]

Point 2. from above does not apply here, since supplying a msg: str overwrites the default diagnostic message.

[NicolasHug]

Thanks Philip, LGTM