diff --git a/models/demos/resnet/tests/test_metal_resnet50.py b/models/demos/resnet/tests/test_metal_resnet50.py index 34f3547c51d..98157dfc96f 100644 --- a/models/demos/resnet/tests/test_metal_resnet50.py +++ b/models/demos/resnet/tests/test_metal_resnet50.py @@ -182,7 +182,7 @@ def run_trace_model(device, tt_image, tt_resnet50): # Compile tt_resnet50(tt_image_res) # Trace - tid = tt_lib.device.BeginTraceCapture(device, 0, 1500000) + tid = tt_lib.device.BeginTraceCapture(device, 0) tt_output_res = tt_resnet50(tt_image_res) tt_lib.device.EndTraceCapture(device, 0, tid) @@ -257,7 +257,7 @@ def run_trace_2cq_model(device, tt_image, tt_resnet50): reshard_out = tt_lib.tensor.reshard(tt_image_res, reshard_mem_config) tt_lib.device.RecordEvent(device, 0, op_event) - tid = tt_lib.device.BeginTraceCapture(device, 0, 1500000) + tid = tt_lib.device.BeginTraceCapture(device, 0) tt_output_res = tt_resnet50(reshard_out, final_out_mem_config=interleaved_dram_mem_config) reshard_out = tt_lib.tensor.allocate_tensor_on_device( reshard_out.shape, reshard_out.dtype, reshard_out.layout, device, reshard_mem_config diff --git a/models/demos/resnet/tests/test_metal_resnet50_2cqs_performant.py b/models/demos/resnet/tests/test_metal_resnet50_2cqs_performant.py index 90af4f781bc..1d98feb58ac 100644 --- a/models/demos/resnet/tests/test_metal_resnet50_2cqs_performant.py +++ b/models/demos/resnet/tests/test_metal_resnet50_2cqs_performant.py @@ -42,7 +42,9 @@ def test_run_resnet50_2cqs_inference( @skip_for_wormhole_b0("This test is not supported on WHB0, please use the TTNN version.") -@pytest.mark.parametrize("device_params", [{"l1_small_size": 24576, "num_hw_cqs": 2}], indirect=True) +@pytest.mark.parametrize( + "device_params", [{"l1_small_size": 24576, "num_hw_cqs": 2, "trace_region_size": 1500000}], indirect=True +) @pytest.mark.parametrize("batch_size", [20], ids=["batch_20"]) @pytest.mark.parametrize( "weights_dtype", diff --git a/models/demos/resnet/tests/test_metal_resnet50_performant.py b/models/demos/resnet/tests/test_metal_resnet50_performant.py index e51b4fb3bd7..535bac8dc77 100644 --- a/models/demos/resnet/tests/test_metal_resnet50_performant.py +++ b/models/demos/resnet/tests/test_metal_resnet50_performant.py @@ -42,7 +42,7 @@ def test_run_resnet50_inference( @skip_for_wormhole_b0("This test is not supported on WHB0, please use the TTNN version.") -@pytest.mark.parametrize("device_params", [{"l1_small_size": 24576}], indirect=True) +@pytest.mark.parametrize("device_params", [{"l1_small_size": 24576, "trace_region_size": 1500000}], indirect=True) @pytest.mark.parametrize("batch_size", [20], ids=["batch_20"]) @pytest.mark.parametrize( "weights_dtype",