29,30c29
< from datasets import load_dataset
< 
---
> import torch
31a31
> from datasets import load_dataset
36d35
<     DataCollatorWithPadding,
40,41d38
<     Trainer,
<     TrainingArguments,
48a46,49
> from optimum.graphcore import IPUConfig, IPUTrainer
> from optimum.graphcore import IPUTrainingArguments as TrainingArguments
> from optimum.graphcore.utils import check_min_version as gc_check_min_version
> 
52a54,56
> # Will error if the minimal version of Optimum Graphcore is not installed. Remove at your own risks.
> gc_check_min_version("0.6.0.dev0")
> 
241,245d244
<     # Log on each process the small summary:
<     logger.warning(
<         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}"
<     )
365a365,370
>     ipu_config = IPUConfig.from_pretrained(
>         training_args.ipu_config_name if training_args.ipu_config_name else model_args.model_name_or_path,
>         cache_dir=model_args.cache_dir,
>         revision=model_args.model_revision,
>         use_auth_token=True if model_args.use_auth_token else None,
>     )
382a388,392
>     # Customize tokenization for GPT2. We reuse the EOS token as the PAD token.
>     if config.model_type == "gpt2":
>         tokenizer.pad_token = tokenizer.eos_token
>         model.config.pad_token_id = model.config.eos_token_id
> 
464a475,483
>         labels = torch.tensor(train_dataset[0]["label"])
>         if model.config.problem_type is None:
>             if model.config.num_labels == 1:
>                 model.config.problem_type = "regression"
>             elif model.config.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
>                 model.config.problem_type = "single_label_classification"
>             else:
>                 model.config.problem_type = "multi_label_classification"
> 
508,509d526
<     elif training_args.fp16:
<         data_collator = DataCollatorWithPadding(tokenizer, pad_to_multiple_of=8)
514c531
<     trainer = Trainer(
---
>     trainer = IPUTrainer(
515a533
>         ipu_config=ipu_config,