Method _assert_close

examples/mlperf/models/test_apply_grad.py:34–36

(self, got, expected, atol, rtol)

    return [p.grad for p in get_parameters(model)]

  def _assert_close(self, got, expected, atol, rtol):
    for g, e in zip(got, expected):
      self.assertTrue(g.allclose(e, atol=atol, rtol=rtol).item(), f"grad mismatch (max abs diff {(g - e).abs().max().item()})")

  def _assert_match(self, model, xs, atol, rtol):
    self._assert_close(self._run_with_apply_grad(model, xs), self._run_reference(model, xs), atol, rtol)