trainer

huggingface · Nov 3, 2023 · 8c77efb · 8c77efb
1 parent 3a706f8
commit 8c77efb
Show file tree

Hide file tree

Showing 2 changed files with 205 additions and 92 deletions.
diff --git a/optimum/intel/neural_compressor/trainer.py b/optimum/intel/neural_compressor/trainer.py
@@ -35,10 +35,9 @@
 from transformers import Trainer
 from transformers.data.data_collator import DataCollator
 from transformers.debug_utils import DebugOption, DebugUnderflowOverflow
-from transformers.file_utils import WEIGHTS_NAME
 
 # Integrations must be imported before ML frameworks:
-from transformers.integrations import deepspeed_init, deepspeed_load_checkpoint, hp_params
+from transformers.integrations import deepspeed_init, deepspeed_load_checkpoint, hp_params, is_deepspeed_available
 from transformers.modeling_utils import PreTrainedModel, get_parameter_dtype, unwrap_model
 from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING_NAMES
 from transformers.pytorch_utils import is_torch_less_than_1_11