commit change

mosaicml · Jun 17, 2024 · 593e6cd · 593e6cd
1 parent 9f00f61
commit 593e6cd
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tests/a_scripts/inference/test_convert_composer_to_hf.py b/tests/a_scripts/inference/test_convert_composer_to_hf.py
@@ -926,9 +926,9 @@ def test_huggingface_conversion_callback(
     # summon full params to check equivalence
     from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 
-    context_manager = te.onnx_export(
+    context_manager = te.onnx_export(  # type: ignore
         True,
-    ) if trainer_precision == 'amp_fp8' else contextlib.nullcontext()  # type: ignore
+    ) if trainer_precision == 'amp_fp8' else contextlib.nullcontext()
     with context_manager:
         with FSDP.summon_full_params(
             trainer.state.model,