fixed SSM sharding test

greg-kwasniewski1 · greg-kwasniewski1 · commit 6149ede4db4b · 2026-01-22T04:23:09.000-08:00
Signed-off-by: greg-kwasniewski1 &lt;213329731+greg-kwasniewski1@users.noreply.github.com&gt;
diff --git a/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_tp_sharding.py b/tests/unittest/_torch/auto_deploy/unit/multigpu/transformations/library/test_tp_sharding.py
@@ -572,7 +572,9 @@ def _run_pattern_detection_job(
                             fused_weight_dims=None,
                         )
                     )
-                if len(node.args) > 1 and "norm_weight" in node.args[0].name:
+                if len(node.args) > 1 and (
+                    "norm_weight" in node.args[0].name or "a_log" in node.args[0].name
+                ):
                     expected_transformations.append(
                         WeightShardingInfo(
                             target_node=node.name,

Original file line number	Diff line number	Diff line change
`@@ -572,7 +572,9 @@ def _run_pattern_detection_job(`
`572`	`572`	`fused_weight_dims=None,`
`573`	`573`	`)`
`574`	`574`	`)`
`575`		`- if len(node.args) > 1 and "norm_weight" in node.args[0].name:`
	`575`	`+ if len(node.args) > 1 and (`
	`576`	`+ "norm_weight" in node.args[0].name or "a_log" in node.args[0].name`
	`577`	`+ ):`
`576`	`578`	`expected_transformations.append(`
`577`	`579`	`WeightShardingInfo(`
`578`	`580`	`target_node=node.name,`