#8156: Add golden function

tenstorrent · Oct 28, 2024 · 1f7e903 · 1f7e903
1 parent 3bdfd9d
commit 1f7e903
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/tests/sweep_framework/sweeps/eltwise/unary/mish.py b/tests/sweep_framework/sweeps/eltwise/unary/mish.py
@@ -6,7 +6,6 @@
 from functools import partial
 
 import torch
-import random
 import ttnn
 from tests.sweep_framework.sweep_utils.utils import gen_shapes
 from tests.tt_eager.python_api_testing.sweep_tests.generation_funcs import gen_func_with_cast_tt
@@ -69,8 +68,8 @@ def run(
         partial(torch_random, low=-100, high=100, dtype=torch.float32), input_a_dtype
     )(input_shape)
 
-    mish = torch.nn.Mish()
-    torch_output_tensor = mish(torch_input_tensor_a)
+    golden_function = ttnn.get_golden_function(ttnn.mish)
+    torch_output_tensor = golden_function(torch_input_tensor_a)
 
     input_tensor_a = ttnn.from_torch(
         torch_input_tensor_a,