debug: add debugging log

build/torch-neuron/__init__.py

@@ -115,6 +115,7 @@ def rmsnorm(hidden_states, weight, eps: float = 1e-6):
 
     # 4. Process in chunks using NKI's hardware-optimized affine_range
     # (Assuming num_rows is perfectly divisible by max_rows for standard tiling)
+    print("Num rows:", num_rows, "Max rows per tile:", max_rows)
     for i in nl.affine_range(num_rows // max_rows):
         
         # Calculate the exact memory offset for this specific chunk