How to use DirectML with transformers TrainingArguments

[osadchi]

Thank you, for your work!!
I have a repo for training llama models, but it's for cuda.
How would I fix it to use DirectML with multi GPU instead of CUDA (I have AMD ROCm and NVIDIA CUDA cards)?

https://github.com/RuslanPeresy/gptchain/blob/main/train.py

There's CUDA everywhere :C

fp16=not torch.cuda.is_bf16_supported(),
bf16=torch.cuda.is_bf16_supported(),

Thank you Again!