Almost exactly two

x = torch.tensor(2.0).sqrt() ** 2

x == 2.0                      # False
x.item()                      # 1.9999998807907104
torch.allclose(x, torch.tensor(2.0))   # True

Nothing is wrong. sqrt(2) is irrational and float32 has 24 bits of mantissa, so squaring the stored approximation lands next door to 2 rather than on it.

What to use instead

torch.allclose(a, b)         # one bool for the whole tensor
torch.isclose(a, b)          # elementwise, a bool tensor

Both test |a - b| <= atol + rtol * |b|. Two tolerances, because the right answer depends on scale:

• rtol (default 1e-5) handles large values, where being within a millionth is as close as float32 gets.
• atol (default 1e-8) handles values near zero, where a relative tolerance would demand impossible precision.

A comparison with only one of them is wrong at one end. That is why there are two, and why atol=0 is a mistake people make when a test is flaky near zero.

Where exact equality is still right

Integers, booleans, and values you put there yourself and never did arithmetic on: a padding sentinel, an index, a class label. Exactness is meaningful for those, and allclose on integer tensors is just slower.

The rule is about arithmetic, not about the dtype: a float that has been computed should be compared with a tolerance, and one that has merely been stored can be compared exactly.

In tests

torch.testing.assert_close is the version to reach for in a test suite. It picks tolerances per dtype (looser for float16 than for float64, which is correct and which you would otherwise get wrong), and its failure message names the worst offending element instead of saying False.

Your task

def close_enough(a: torch.Tensor, b: torch.Tensor, atol: float, rtol: float) -> bool

Return whether every element of a is within tolerance of b, using the given atol and rtol. Return a Python bool.