From 3aad788327a3b0b813e6807fb9228966c031b143 Mon Sep 17 00:00:00 2001
From: Borys Bradel <164946524+bbradelTT@users.noreply.github.com>
Date: Fri, 13 Sep 2024 12:39:48 -0400
Subject: [PATCH] #12619: Update matmul sweep timeout and core range set usage
 (#12655)

---
 .../sweeps/matmul/full/matmul_default_height_sharded.py         | 2 +-
 .../sweeps/matmul/full/matmul_default_width_sharded.py          | 2 +-
 tests/sweep_framework/sweeps/matmul/short/matmul_default.py     | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/sweep_framework/sweeps/matmul/full/matmul_default_height_sharded.py b/tests/sweep_framework/sweeps/matmul/full/matmul_default_height_sharded.py
index 44cd92fe128..fff716bc2ea 100644
--- a/tests/sweep_framework/sweeps/matmul/full/matmul_default_height_sharded.py
+++ b/tests/sweep_framework/sweeps/matmul/full/matmul_default_height_sharded.py
@@ -88,7 +88,7 @@ def run(
     # TODO: row_wise=False and ROW_MAJOR shard orientation gives bad PCC
     # TODO: COL_MAJOR shard orientation doesn't work for get_matmul_program_config
     input_a_memory_config.shard_spec = ttnn.ShardSpec(
-        ttnn.num_cores_to_corerange_set(num_cores_height, core_grid, row_wise=True),
+        ttnn.CoreRangeSet(ttnn.num_cores_to_corerange_set(num_cores_height, core_grid, row_wise=True)),
         (per_core_height, k_size),
         ttnn.ShardOrientation.ROW_MAJOR,
         False,
diff --git a/tests/sweep_framework/sweeps/matmul/full/matmul_default_width_sharded.py b/tests/sweep_framework/sweeps/matmul/full/matmul_default_width_sharded.py
index e6fb7562639..a7f0361f4fe 100644
--- a/tests/sweep_framework/sweeps/matmul/full/matmul_default_width_sharded.py
+++ b/tests/sweep_framework/sweeps/matmul/full/matmul_default_width_sharded.py
@@ -81,7 +81,7 @@ def run(
     # TODO: row_wise=False and ROW_MAJOR shard orientation gives bad PCC
     # TODO: COL_MAJOR shard orientation doesn't work for get_matmul_program_config
     input_a_memory_config.shard_spec = ttnn.ShardSpec(
-        ttnn.num_cores_to_corerange_set(num_cores_width, core_grid, row_wise=True),
+        ttnn.CoreRangeSet(ttnn.num_cores_to_corerange_set(num_cores_width, core_grid, row_wise=True)),
         (total_height, per_core_width),
         ttnn.ShardOrientation.ROW_MAJOR,
         False,
diff --git a/tests/sweep_framework/sweeps/matmul/short/matmul_default.py b/tests/sweep_framework/sweeps/matmul/short/matmul_default.py
index d872fcf1874..bb12db78c58 100644
--- a/tests/sweep_framework/sweeps/matmul/short/matmul_default.py
+++ b/tests/sweep_framework/sweeps/matmul/short/matmul_default.py
@@ -12,7 +12,7 @@
 from models.utility_functions import torch_random
 
 
-TIMEOUT = 5
+TIMEOUT = 15
 
 # TODO: Missing coverage for mixed precision; passed in dtype does nothing in current matmul path
 parameters = {