Merge pull request #779 from JiabinYang/change_device

fix change device error code
deepmodeling · Jun 22, 2021 · 55670b2 · 55670b2
2 parents 4b24e1f + 9a92b7f
commit 55670b2
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/deepmd/train/trainer.py b/deepmd/train/trainer.py
@@ -80,6 +80,7 @@ class DPTrainer (object):
     def __init__(self, 
                  jdata, 
                  run_opt):
+        paddle.set_device("cpu")
         self.run_opt = run_opt
         self._init_param(jdata)
 
@@ -304,7 +305,7 @@ def build (self,
     def train (self, 
                data,
                stop_batch) :
-        paddle.set_device("gpu")
+
         self.stop_batch = stop_batch
 
         self.print_head()