Browse Source

ts models only support single gpu to train, eval, etc.

gaotingquan 1 year ago
parent
commit
09e80570ef

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

@@ -117,5 +117,5 @@ exporting!"
         return {
             "weight_path": self.export_config.weight_path,
             "save_dir": self.global_config.output,
-            "device": self.get_device(),
+            "device": self.get_device(1),
         }

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

@@ -91,7 +91,7 @@ training!"
         Returns:
             dict: the arguments of training function.
         """
-        train_args = {"device": self.get_device()}
+        train_args = {"device": self.get_device(using_device_number=1)}
         if self.global_config.output is not None:
             train_args["save_dir"] = self.global_config.output
         return train_args

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

@@ -86,7 +86,7 @@ training!"
         Returns:
             dict: the arguments of training function.
         """
-        train_args = {"device": self.get_device()}
+        train_args = {"device": self.get_device(using_device_number=1)}
         if self.global_config.output is not None:
             train_args["save_dir"] = self.global_config.output
         return train_args

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

@@ -86,7 +86,7 @@ training!"
         Returns:
             dict: the arguments of training function.
         """
-        train_args = {"device": self.get_device()}
+        train_args = {"device": self.get_device(using_device_number=1)}
         if self.global_config.output is not None:
             train_args["save_dir"] = self.global_config.output
         return train_args