Explorar el Código

fix env setting not work when using single gpu

gaotingquan hace 1 año
padre
commit
e8bfa2e532

+ 1 - 1
paddlex/modules/base/evaluator.py

@@ -138,9 +138,9 @@ evaling!"
         Returns:
             str: device setting, such as: `gpu:0,1`, `npu:0,1`, `cpu`.
         """
+        set_env_for_device(self.global_config.device)
         if using_device_number:
             return update_device_num(self.global_config.device, using_device_number)
-        set_env_for_device(self.global_config.device)
         return self.global_config.device
 
     @abstractmethod

+ 1 - 1
paddlex/modules/base/exportor.py

@@ -103,9 +103,9 @@ exporting!"
         Returns:
             str: device setting, such as: `gpu:0,1`, `npu:0,1`, `cpu`.
         """
+        set_env_for_device(self.global_config.device)
         if using_device_number:
             return update_device_num(self.global_config.device, using_device_number)
-        set_env_for_device(self.global_config.device)
         return self.global_config.device
 
     def update_config(self):

+ 1 - 1
paddlex/modules/base/trainer.py

@@ -95,9 +95,9 @@ training!"
         Returns:
             str: device setting, such as: `gpu:0,1`, `npu:0,1` `cpu`.
         """
+        set_env_for_device(self.global_config.device)
         if using_device_number:
             return update_device_num(self.global_config.device, using_device_number)
-        set_env_for_device(self.global_config.device)
         return self.global_config.device
 
     @abstractmethod