fix when tensor include nan (#5724)(#5725)

mayuyuace · web-flow · commit dd525e87e40c · 2025-07-11T16:38:51.000+08:00
diff --git a/tests/gpu/examples/moe/test_moe_gemm.py b/tests/gpu/examples/moe/test_moe_gemm.py
@@ -95,5 +95,9 @@ def test_moe_gemm(
                 n_experts,
             )
             torch.testing.assert_close(
-                output_fp8.to(float), ref_output.to(float), rtol=1e-2, atol=1e-2
+                output_fp8.to(float),
+                ref_output.to(float),
+                rtol=1e-2,
+                atol=1e-2,
+                equal_nan=True,
             )

Original file line number	Diff line number	Diff line change
`@@ -95,5 +95,9 @@ def test_moe_gemm(`
`95`	`95`	`n_experts,`
`96`	`96`	`)`
`97`	`97`	`torch.testing.assert_close(`
`98`		`- output_fp8.to(float), ref_output.to(float), rtol=1e-2, atol=1e-2`
	`98`	`+ output_fp8.to(float),`
	`99`	`+ ref_output.to(float),`
	`100`	`+ rtol=1e-2,`
	`101`	`+ atol=1e-2,`
	`102`	`+ equal_nan=True,`
`99`	`103`	`)`