Merge pull request #1593 from windstamp/npu_dev_20220322

[NPU] Add NPU support for TransformerTTS
PaddlePaddle · Mar 23, 2022 · 26ef478 · 26ef478
2 parents 6813002 + 59b3de6
commit 26ef478
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/paddlespeech/t2s/exps/transformer_tts/train.py b/paddlespeech/t2s/exps/transformer_tts/train.py
@@ -42,10 +42,12 @@
 def train_sp(args, config):
     # decides device type and whether to run in parallel
     # setup running environment correctly
-    if (not paddle.is_compiled_with_cuda()) or args.ngpu == 0:
-        paddle.set_device("cpu")
-    else:
+    if paddle.is_compiled_with_cuda() and args.ngpu > 0:
         paddle.set_device("gpu")
+    elif paddle.is_compiled_with_npu() and args.ngpu > 0:
+        paddle.set_device("npu")
+    else:
+        paddle.set_device("cpu")
     world_size = paddle.distributed.get_world_size()
     if world_size > 1:
         paddle.distributed.init_parallel_env()