diff --git a/rl_engine/executors/deepspeed_trainer.py b/rl_engine/executors/deepspeed_trainer.py
index 9eb276c..fa0860b 100644
--- a/rl_engine/executors/deepspeed_trainer.py
+++ b/rl_engine/executors/deepspeed_trainer.py
@@ -490,6 +490,7 @@ def _coerce_hidden_tensor(
                 return hidden
     if hasattr(candidate, "hidden_states"):
         hidden = _last_hidden_state_tensor(
+            getattr(candidate, "hidden_states"),
             candidate.hidden_states,
             expected_hidden_dim=expected_hidden_dim,
         )
diff --git a/tests/test_linear_logp.py b/tests/test_linear_logp.py
index 659127c..bdd6cd6 100644
--- a/tests/test_linear_logp.py
+++ b/tests/test_linear_logp.py
@@ -312,63 +312,6 @@ def test_chunked_linear_logp_backward_matches_autograd_and_layout_invariance(use
                 assert torch.allclose(grad_bias, canonical_bias_grad, atol=1e-6)
 
 
-def test_chunked_linear_logp_backward_skips_unused_gradients():
-    hidden, weight, target, bias = _inputs(2029, device="cpu", bias=True)
-    grad_out = torch.randn_like(target, dtype=torch.float32)
-    hidden_2d = hidden.reshape(-1, hidden.size(-1)).contiguous()
-    target_1d = target.reshape(-1).contiguous()
-
-    full_hidden, full_weight, full_bias = chunked_linear_logp_backward(
-        grad_out,
-        hidden_2d,
-        weight,
-        target_1d,
-        bias,
-        has_bias=True,
-        lead_shape=target.shape,
-        hidden_dtype=hidden.dtype,
-        weight_dtype=weight.dtype,
-        bias_dtype=bias.dtype,
-        chunk_elems=weight.size(0) * 3,
-    )
-    skipped_hidden, skipped_weight, skipped_bias = chunked_linear_logp_backward(
-        grad_out,
-        hidden_2d,
-        weight,
-        target_1d,
-        bias,
-        has_bias=True,
-        lead_shape=target.shape,
-        hidden_dtype=hidden.dtype,
-        weight_dtype=weight.dtype,
-        bias_dtype=bias.dtype,
-        chunk_elems=weight.size(0) * 3,
-        compute_grad_hidden=False,
-    )
-    only_hidden, skipped_weight_2, skipped_bias_2 = chunked_linear_logp_backward(
-        grad_out,
-        hidden_2d,
-        weight,
-        target_1d,
-        bias,
-        has_bias=True,
-        lead_shape=target.shape,
-        hidden_dtype=hidden.dtype,
-        weight_dtype=weight.dtype,
-        bias_dtype=bias.dtype,
-        chunk_elems=weight.size(0) * 3,
-        compute_grad_weight=False,
-        compute_grad_bias=False,
-    )
-
-    assert skipped_hidden is None
-    assert torch.allclose(skipped_weight, full_weight, atol=1e-6)
-    assert torch.allclose(skipped_bias, full_bias, atol=1e-6)
-    assert skipped_weight_2 is None
-    assert skipped_bias_2 is None
-    assert torch.allclose(only_hidden, full_hidden, atol=1e-6)
-
-
 def test_tied_embedding_lm_head_shared_gradient_is_layout_invariant():
     torch.manual_seed(2028)
     model = _EmbeddingLMHeadModel(vocab_size=13, hidden_dim=6, bias=False, tie_weights=True)