wenet icon indicating copy to clipboard operation
wenet copied to clipboard

[device args support] Replace `gpu` args with `device` args in train.py/recognize.py, etc.

Open MengqingCao opened this issue 1 year ago • 0 comments

Is your feature request related to a problem? Please describe. I'm running wenet on Ascend NPU device recently. And I find that the scripts, e.g., train.py, recognize.py, ./cli/modle.py, only support GPU device via --gpu argument.

Describe the solution you'd like Specifying the accelerator to be used through --device argument will make it more convenient for people using wenet with different devices. There are a lot of devices (gpu, mps, xpu, etc.) supported in PyTorch, including Ascend NPU. If the replacement brings too many changes, adding a --device argumentcould also solve the problem, although I think the way of replacing is more reasonable.

Please let me know if there is any suggestions or concerns, thx!

MengqingCao avatar May 06 '24 08:05 MengqingCao