20d19
< import warnings
28,29d26
< from datasets import DatasetDict, load_dataset
< 
31,39c28,29
< from transformers import (
<     AutoConfig,
<     AutoFeatureExtractor,
<     AutoModelForAudioClassification,
<     HfArgumentParser,
<     Trainer,
<     TrainingArguments,
<     set_seed,
< )
---
> from datasets import DatasetDict, load_dataset
> from transformers import AutoConfig, AutoFeatureExtractor, AutoModelForAudioClassification, HfArgumentParser
43a34,36
> from optimum.habana import GaudiConfig, GaudiTrainer, GaudiTrainingArguments
> from optimum.habana.utils import set_seed
> 
48c41
< check_min_version("4.29.0.dev0")
---
> check_min_version("4.28.0")
164,166d156
<     freeze_feature_extractor: Optional[bool] = field(
<         default=None, metadata={"help": "Whether to freeze the feature extractor layers of the model."}
<     )
172,186d161
<     def __post_init__(self):
<         if not self.freeze_feature_extractor and self.freeze_feature_encoder:
<             warnings.warn(
<                 "The argument `--freeze_feature_extractor` is deprecated and "
<                 "will be removed in a future version. Use `--freeze_feature_encoder`"
<                 "instead. Setting `freeze_feature_encoder==True`.",
<                 FutureWarning,
<             )
<         if self.freeze_feature_extractor and not self.freeze_feature_encoder:
<             raise ValueError(
<                 "The argument `--freeze_feature_extractor` is deprecated and "
<                 "should not be used in combination with `--freeze_feature_encoder`."
<                 "Only make use of `--freeze_feature_encoder`."
<             )
< 
193c168
<     parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
---
>     parser = HfArgumentParser((ModelArguments, DataTrainingArguments, GaudiTrainingArguments))
221a197,203
>     gaudi_config = GaudiConfig.from_pretrained(
>         training_args.gaudi_config_name,
>         cache_dir=model_args.cache_dir,
>         revision=model_args.model_revision,
>         use_auth_token=True if model_args.use_auth_token else None,
>     )
> 
224,225c206,208
<         f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu} "
<         + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
---
>         f"Process rank: {training_args.local_rank}, device: {training_args.device}, "
>         + f"distributed training: {bool(training_args.local_rank != -1)}, "
>         + f"mixed-precision training: {gaudi_config.use_habana_mixed_precision}"
291a275,277
>     # Max input length
>     max_length = int(round(feature_extractor.sampling_rate * data_args.max_length_seconds))
> 
296a283
> 
302c289,295
<         inputs = feature_extractor(subsampled_wavs, sampling_rate=feature_extractor.sampling_rate)
---
>         inputs = feature_extractor(
>             subsampled_wavs,
>             max_length=max_length,
>             sampling_rate=feature_extractor.sampling_rate,
>             padding="max_length",
>             truncation=True,
>         )
311c304,310
<         inputs = feature_extractor(wavs, sampling_rate=feature_extractor.sampling_rate)
---
>         inputs = feature_extractor(
>             wavs,
>             max_length=max_length,
>             sampling_rate=feature_extractor.sampling_rate,
>             padding="max_length",
>             truncation=True,
>         )
376c375
<     trainer = Trainer(
---
>     trainer = GaudiTrainer(
377a377
>         gaudi_config=gaudi_config,
