fix tests

huggingface · Dec 15, 2023 · e8b3cc5 · e8b3cc5
1 parent 8b6b801
commit e8b3cc5
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/fsdp/test_fsdp.py b/tests/fsdp/test_fsdp.py
@@ -34,14 +34,14 @@
     slow,
     torch_device,
 )
-from transformers.trainer import FSDP_MODEL_NAME
 from transformers.trainer_callback import TrainerState
 from transformers.trainer_utils import FSDPOption, set_seed
 from transformers.utils import is_accelerate_available, is_torch_bf16_available_on_device
 
 
 if is_torch_available():
     from transformers.pytorch_utils import is_torch_greater_or_equal_than_2_1
+    from transformers.trainer import FSDP_MODEL_NAME
 else:
     is_torch_greater_or_equal_than_2_1 = False