diff --git a/tensorrt_llm/_torch/models/modeling_nemotron_h.py b/tensorrt_llm/_torch/models/modeling_nemotron_h.py
index 249cab1547..d1a4974e74 100644
--- a/tensorrt_llm/_torch/models/modeling_nemotron_h.py
+++ b/tensorrt_llm/_torch/models/modeling_nemotron_h.py
@@ -361,6 +361,10 @@ class NemotronHLayer(DecoderLayer):
                                    spec_metadata=spec_metadata,
                                    **kwargs)
         hidden_states = torch.add(hidden_states, residual)
+        if spec_metadata is not None and spec_metadata.is_layer_capture(
+                self.layer_idx):
+            spec_metadata.maybe_capture_hidden_states(self.layer_idx,
+                                                      hidden_states, None)
 
         return hidden_states