fix nemotron capture for non attention layers (#21436)

2026-03-30 12:50:49 -07:00
parent ad064c2f4e
commit 4a9ffc3ab6
1 changed files with 2 additions and 0 deletions
@@ -2424,6 +2424,8 @@ class ModelRunner(ModelRunnerKVCacheMixin):

            if attn_layer is not None:
                self.attention_layers.append(attn_layer)
+            elif hasattr(layer, "mixer"):
+                self.attention_layers.append(None)

            moe_block = None
            moe_fusion = None