From 24874f90a3246a7c2895cea13cc0e4cdb875280e Mon Sep 17 00:00:00 2001 From: YAMY <74099316+YAMY1234@users.noreply.github.com> Date: Tue, 15 Sep 2026 12:40:28 -0700 Subject: [PATCH] [Test] Fix DeepGEMM batch invariance test output dtype (#39636) --- .../unit/batch_invariant_ops/test_batch_invariant_ops.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/test/registered/unit/batch_invariant_ops/test_batch_invariant_ops.py b/test/registered/unit/batch_invariant_ops/test_batch_invariant_ops.py index 21a11142c..a89b64d5e 100644 --- a/test/registered/unit/batch_invariant_ops/test_batch_invariant_ops.py +++ b/test/registered/unit/batch_invariant_ops/test_batch_invariant_ops.py @@ -117,10 +117,12 @@ class TestBatchInvariantOps(CustomTestCase): if mode is not None: deep_gemm.use_deterministic_algorithms(mode) with self.subTest(deterministic=mode): - ref = batch_invariant_ops._matmul_persistent_deepgemm(a[:1], b) + ref = batch_invariant_ops._matmul_persistent_deepgemm( + a[:1], b, out_dtype=torch.bfloat16 + ) for batch_size in (16, 64, 257): out = batch_invariant_ops._matmul_persistent_deepgemm( - a[:batch_size], b + a[:batch_size], b, out_dtype=torch.bfloat16 ) torch.testing.assert_close(out[:1], ref, rtol=0, atol=0) if get_deterministic: