fixed tanhshrink cuda bfloat16 failure

parth-desai · parth-desai · commit 97b9196ca0bb · 2023-08-29T19:36:53.000-07:00
diff --git a/torch/testing/_internal/common_methods_invocations.py b/torch/testing/_internal/common_methods_invocations.py
@@ -13739,7 +13739,8 @@ def reference_flatten(input, start_dim=0, end_dim=-1):
                 toleranceOverride({torch.bfloat16: tol(atol=1e-02, rtol=1.6e-02)}), 'TestUnaryUfuncs',),
             DecorateInfo(toleranceOverride({torch.float16: tol(atol=1e-02, rtol=1e-02)}),
                          'TestCudaFuserOpInfo', 'test_nvfuser_correctness'),
-            DecorateInfo(toleranceOverride({torch.complex64: tol(atol=6e-04, rtol=1e-05)}),
+            DecorateInfo(toleranceOverride({torch.complex64: tol(atol=6e-04, rtol=1e-05),
+                                            torch.bfloat16: tol(atol=1e-02, rtol=1.6e-02)}),
                          'TestUnaryUfuncs', 'test_reference_numerics_extremal',
                          device_type='cuda'),
         ],
