fixing comments from AI

Signed-off-by: Govind Ramnarayan <105831528+govind-ramnarayan@users.noreply.github.com>

Author: govind-ramnarayan

tensorrt_llm/_torch/auto_deploy/llm_args.py

@@ -46,7 +46,7 @@ def _check_for_default_value_only(
 
 
 def default_eagle3_layers_to_capture(num_hidden_layers: int) -> Set[int]:
-    if num_hidden_layers <= 5:
+    if num_hidden_layers <= 6:
         raise ValueError("Not enough hidden layers for default EAGLE3 capture")
     return {1, num_hidden_layers // 2 - 1, num_hidden_layers - 4}
 
@@ -415,7 +415,7 @@ def ensure_no_custom_parallel_config(cls, value: Any, info: ValidationInfo) -> A
         return _check_for_default_value_only(cls, value, info, msg)
 
     @model_validator(mode="after")
-    def default_eagle3_layers_to_capture(self):
+    def set_eagle3_layers_to_capture(self):
         if self.speculative_config is None or not isinstance(
             self.speculative_config, EagleDecodingConfig
         ):
@@ -428,6 +428,12 @@ def default_eagle3_layers_to_capture(self):
             )
 
         # insert the layers to capture into the transforms config.
+        if self.transforms is None:
+            self.transforms = {}
+
+        if "detect_hidden_states_for_capture" not in self.transforms:
+            self.transforms["detect_hidden_states_for_capture"] = {}
+
         self.transforms["detect_hidden_states_for_capture"]["eagle3_layers_to_capture"] = (
             self.speculative_config.eagle3_layers_to_capture
         )

tests/integration/defs/examples/test_ad_speculative_decoding.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 
 import os
+from typing import Optional
 
 import pytest
 from build_and_run_ad import ExperimentConfig, main
@@ -65,7 +66,9 @@ def make_spec_config(spec_dec_mode: str, spec_model_path: str):
     raise ValueError(f"Unknown speculative mode: {spec_dec_mode}")
 
 
-def run_with_autodeploy(model, speculative_model_dir, batch_size, spec_dec_mode: str | None):
+def run_with_autodeploy(
+    model, speculative_model_dir, batch_size, spec_dec_mode: Optional[str] = None
+):
     """Run AutoDeploy with or without speculative decoding.
 
     Args: