NameError: name 'PartialState' is not defined error while training hugging face wave2vec model

Here is the code block which caused the error

training_args = TrainingArguments(
    output_dir="my_awesome_mind_model",
    evaluation_strategy="epoch",
    save_strategy="epoch",
    learning_rate=3e-5,
    per_device_train_batch_size=32,
    gradient_accumulation_steps=4,
    per_device_eval_batch_size=32,
    num_train_epochs=10,
    warmup_ratio=0.1,
    logging_steps=10,
    load_best_model_at_end=True,
    metric_for_best_model="accuracy",
    push_to_hub=True,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset["train"],
    # eval_dataset=encoded_minds["test"],
    tokenizer=feature_extractor,
    compute_metrics=compute_metrics,
)

trainer.train()

getting the following error

NameError                                 Traceback (most recent call last)
 in <cell line: 1>()
1 training_args = TrainingArguments(
2     output_dir="my_awesome_mind_model",
3     evaluation_strategy="epoch",
4     save_strategy="epoch",
5     learning_rate=3e-5,

4 frames
/usr/local/lib/python3.10/dist-packages/transformers/training_args.py in _setup_devices(self)
1629             self._n_gpu = 1
1630         else:
1631             self.distributed_state = PartialState(backend=self.ddp_backend)
1632             self._n_gpu = 1
1633         if not is_sagemaker_mp_enabled():

NameError: name 'PartialState' is not defined

I am trying to follow the audio classification guide of hugging face(link on another dataset but upon running the training args code i am getting name "PartialState" not defined error.

Get this solution programmatically \u2014 free, no authentication.