Fix HF model dtype fallback in forward_pass_logit_checker

subawocit · subawocit · commit aecb57d511da · 2026-06-30T22:58:40.000Z
diff --git a/tests/utils/forward_pass_logit_checker.py b/tests/utils/forward_pass_logit_checker.py
@@ -428,7 +428,7 @@ def main(config, test_args):  # pylint: disable=W0621
     }
 
     # Default to bfloat16 if dtype is unrecognized
-    torch_dtype = dtype_mapping.get(config.dtype.name, torch.bfloat16)
+    torch_dtype = dtype_mapping.get(config.dtype.name.lower(), torch.bfloat16)
     max_logging.log(f"Loading HF model with dtype: {torch_dtype} (derived from config.dtype: {config.dtype})")
 
     hf_model = AutoModelForCausalLM.from_pretrained(

Original file line number	Diff line number	Diff line change
`@@ -428,7 +428,7 @@ def main(config, test_args): # pylint: disable=W0621`
`428`	`428`	`}`
`429`	`429`
`430`	`430`	`# Default to bfloat16 if dtype is unrecognized`
`431`		`- torch_dtype = dtype_mapping.get(config.dtype.name, torch.bfloat16)`
	`431`	`+ torch_dtype = dtype_mapping.get(config.dtype.name.lower(), torch.bfloat16)`
`432`	`432`	`max_logging.log(f"Loading HF model with dtype: {torch_dtype} (derived from config.dtype: {config.dtype})")`
`433`	`433`
`434`	`434`	`hf_model = AutoModelForCausalLM.from_pretrained(`