pip install torch torchvision torchaudio
[TrainingArguments] uses the mps device by default if it's available which means you don't need to explicitly set the device.