Add cuda check in Transformer. (#2728)

PaddleNLP/neural_machine_translation/transformer/infer.py

@@ -4,12 +4,14 @@ import multiprocessing
 import numpy as np
 import os
 import sys
+sys.path.append("../../")
 sys.path.append("../../models/neural_machine_translation/transformer/")
 from functools import partial
 
 import paddle
 import paddle.fluid as fluid
 
+from models.model_check import check_cuda
 import reader
 from config import *
 from desc import *
@@ -217,6 +219,7 @@ def fast_infer(args):
         fluid.memory_optimize(infer_program)
 
     if InferTaskConfig.use_gpu:
+        check_cuda(InferTaskConfig.use_gpu)
         place = fluid.CUDAPlace(0)
         dev_count = fluid.core.get_cuda_device_count()
     else:

PaddleNLP/neural_machine_translation/transformer/train.py

@@ -6,12 +6,14 @@ import multiprocessing
 import os
 import six
 import sys
+sys.path.append("../../")
 sys.path.append("../../models/neural_machine_translation/transformer/")
 import time
 
 import numpy as np
 import paddle.fluid as fluid
 
+from models.model_check import check_cuda
 import reader
 from config import *
 from desc import *
@@ -663,6 +665,7 @@ def train(args):
         place = fluid.CPUPlace()
         dev_count = int(os.environ.get('CPU_NUM', multiprocessing.cpu_count()))
     else:
+        check_cuda(TrainTaskConfig.use_gpu)
         gpu_id = int(os.environ.get('FLAGS_selected_gpus', 0))
         place = fluid.CUDAPlace(gpu_id)
         dev_count = get_device_num()