Update _utils.py

This commit is contained in:
Daniel Han 2025-10-15 14:30:50 -07:00
commit e657b21cc7

View file

@ -1711,8 +1711,7 @@ def _prepare_model_for_qat(model: torch.nn.Module, qat_scheme: str) -> torch.nn.
while hasttr(model, "model"):
model._torchao_metadata = torchao_metadata
model = model.model
if hasattr(model, "model"):
model._torchao_metadata = torchao_metadata
model._torchao_metadata = torchao_metadata
quantize_(model, QATConfig(base_config, step="prepare"), filter_fn=filter_fn)
return model
pass