Check cases in docs/operations/aten.lt.Scalar.md (#624)

Author: jdh8

tests/lowering/eltwise/binary/test_lt.py

@@ -41,11 +41,23 @@ def test_lt_tensor(device, input_shapes):
     assert torch.allclose(result_before, result_after.to(torch.bool))
 
 
-@pytest.mark.parametrize("input_shape", ((64, 128), (1, 1, 256), (2, 377, 355)))
+@pytest.mark.parametrize(
+    "input_shape",
+    (
+        (64, 128),
+        (1, 1, 256),
+        (2, 377, 355),
+        (1, 1),
+        (10, 10),
+        (15, 15),
+        (17, 17),
+        (2, 2),
+    ),
+)
 def test_lt_scalar(device, input_shape):
     m = LtModule()
     input = torch.randint(0, 10, input_shape, dtype=torch.bfloat16)
-    scalar = torch.randint(0, 10, ()).item()
+    scalar = 5
     result_before = m.forward(input, scalar)
     option = torch_ttnn.TorchTtnnOption(device=device)
     option.gen_graphviz = True
@@ -55,8 +67,9 @@ def test_lt_scalar(device, input_shape):
     option._out_fx_graphs[0].print_tabular()
 
     # Check the graph has be rewritten and contain ttnn ops
-    nodes = list(option._out_fx_graphs[0].nodes)
-    assert [node.target for node in nodes].count(ttnn.lt) == 1
+    nodes = [node.target for node in option._out_fx_graphs[0].nodes]
+    assert torch.ops.aten.lt.Scalar not in nodes
+    assert nodes.count(ttnn.lt) == 1
 
     # Check inference result
     assert torch.allclose(result_before, result_after.to(torch.bool))