18d17
< 
32a32
> import transformers
38,39d37
< 
< import transformers
45,47d42
<     Trainer,
<     TrainingArguments,
<     set_seed,
52a48,50
> from optimum.habana import GaudiConfig, GaudiTrainer, GaudiTrainingArguments
> from optimum.habana.utils import set_seed
> 
57c55
< check_min_version("4.29.0.dev0")
---
> check_min_version("4.28.0")
230c228
<     parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
---
>     parser = HfArgumentParser((ModelArguments, DataTrainingArguments, GaudiTrainingArguments))
258a257,263
>     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,
>     )
> 
261,262c266,268
<         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}"
399d404
<     image_transformations = torch.jit.script(image_transformations)
496c501
<     trainer = Trainer(
---
>     trainer = GaudiTrainer(
497a503
>         gaudi_config=gaudi_config,
