From 6bad7e9e9e9ebe7da2e7ac03e8ea7404fed831c4 Mon Sep 17 00:00:00 2001 From: whn09 Date: Wed, 22 Apr 2026 11:41:19 +0800 Subject: [PATCH] Fix NameError in HuggingFaceGenerationAdapter.prepare_inputs_for_generation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit prepare_inputs_for_generation builds model_inputs with `"tensor_capture_hook": tensor_capture_hook` but never extracts tensor_capture_hook from **kwargs (while input_capture_hook is extracted in the line directly above it). Any call path that reaches this method — including adapter.generate() — raises NameError. The fix is to extract tensor_capture_hook from kwargs, mirroring how input_capture_hook is handled. No other change is needed because NeuronBaseForImageToText.forward already accepts tensor_capture_hook as a keyword argument. Co-Authored-By: Claude Opus 4.7 --- src/neuronx_distributed_inference/utils/hf_adapter.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/neuronx_distributed_inference/utils/hf_adapter.py b/src/neuronx_distributed_inference/utils/hf_adapter.py index 34a2b852..6b81b5b4 100644 --- a/src/neuronx_distributed_inference/utils/hf_adapter.py +++ b/src/neuronx_distributed_inference/utils/hf_adapter.py @@ -278,6 +278,7 @@ def prepare_inputs_for_generation( scatter_index = kwargs.get("scatter_index", None) position_ids = kwargs.get("position_ids", None) input_capture_hook = kwargs.get("input_capture_hook", None) + tensor_capture_hook = kwargs.get("tensor_capture_hook", None) if attention_mask is not None and position_ids is None: # create position_ids on the fly for batch generation