#8282: fix moreh_nllloss callback

tenstorrent · May 27, 2024 · 9f2fd27 · 9f2fd27
1 parent 31282ff
commit 9f2fd27
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 13 deletions.
diff --git a/tests/tt_eager/python_api_testing/unit_testing/misc/test_moreh_nll_loss.py b/tests/tt_eager/python_api_testing/unit_testing/misc/test_moreh_nll_loss.py
@@ -248,9 +248,7 @@ def test_moreh_nll_loss(shape, ignore_index, reduction, none_weight, fp32_dest_a
 @pytest.mark.parametrize("ignore_index", [-1])
 @pytest.mark.parametrize("reduction", ["mean", "sum"])
 @pytest.mark.parametrize("none_weight", [True, False])
-def test_moreh_nll_loss_cache(shape, ignore_index, reduction, none_weight, device, use_program_cache):
-    device.enable_program_cache()
-
+def test_moreh_nll_loss_callback(shape, ignore_index, reduction, none_weight, device, use_program_cache):
     (torch_input, torch_target, torch_weight, torch_divisor, torch_output) = get_torch_tensors(shape)
 
     if none_weight:
@@ -264,15 +262,16 @@ def test_moreh_nll_loss_cache(shape, ignore_index, reduction, none_weight, devic
     )
 
     reduction_mean = reduction == "mean"
-    tt_loss = ttl.operations.primary.moreh_nll_loss(
-        tt_input,
-        tt_target,
-        tt_weight,
-        tt_divisor,
-        tt_output,
-        ignore_index,
-        reduction_mean,
-    )
+    for _ in range(2):
+        tt_loss = ttl.operations.primary.moreh_nll_loss(
+            tt_input,
+            tt_target,
+            tt_weight,
+            tt_divisor,
+            tt_output,
+            ignore_index,
+            reduction_mean,
+        )
 
     tt_loss_to_cpu = tt_loss.cpu().to(ttl.tensor.Layout.ROW_MAJOR).unpad_from_tile([1, 1]).to_torch().reshape([1])
     rtol = atol = 0.05

diff --git a/tt_eager/tt_dnn/op_library/moreh_nll_loss/moreh_nll_loss_step1/moreh_nll_loss_step1.cpp b/tt_eager/tt_dnn/op_library/moreh_nll_loss/moreh_nll_loss_step1/moreh_nll_loss_step1.cpp
@@ -164,7 +164,9 @@ operation::ProgramWithCallbacks moreh_nll_loss_step1_impl(
                 {
                     auto &runtime_args = GetRuntimeArgs(program, reader_kernel_id, core);
                     runtime_args[0] = target_dram_buffer->address();
-                    runtime_args[1] = weight_dram_buffer->address();
+                    if (weight_dram_buffer != nullptr) {
+                        runtime_args[1] = weight_dram_buffer->address();
+                    }
                 }
 
                 {