use default device of platform instead of gpu (#1080)

Author: HydrogenSulfate

ppsci/solver/solver.py

@@ -83,7 +83,7 @@ class Solver:
         use_wandb (Optional[bool]): Whether use wandb to log data. Defaults to False.
         use_tbd (Optional[bool]): Whether use tensorboardX to log data. Defaults to False.
         wandb_config (Optional[Dict[str, str]]): Config dict of WandB. Defaults to None.
-        device (Literal["cpu", "gpu", "xpu"], optional): Runtime device. Defaults to "gpu".
+        device (Literal["cpu", "gpu", "xpu", None], optional): Runtime device. Defaults to None, which means use default device on current platform.
         equation (Optional[Dict[str, ppsci.equation.PDE]]): Equation dict. Defaults to None.
         geom (Optional[Dict[str, ppsci.geometry.Geometry]]): Geometry dict. Defaults to None.
         validator (Optional[Dict[str, ppsci.validate.Validator]]): Validator dict. Defaults to None.
@@ -145,7 +145,7 @@ def __init__(
         use_wandb: bool = False,
         use_tbd: bool = False,
         wandb_config: Optional[Mapping] = None,
-        device: Literal["cpu", "gpu", "xpu"] = "gpu",
+        device: Literal["cpu", "gpu", "xpu", None] = None,
         equation: Optional[Dict[str, ppsci.equation.PDE]] = None,
         geom: Optional[Dict[str, ppsci.geometry.Geometry]] = None,
         validator: Optional[Dict[str, ppsci.validate.Validator]] = None,
@@ -247,7 +247,12 @@ def __init__(
         # set running device
         if not cfg:
             self.device = device
+        if self.device is None:
+            # set to default device if not specified
+            self.device: str = paddle.device.get_device()
+
         if self.device != "cpu" and paddle.device.get_device() == "cpu":
+            # fall back to cpu if no other device available
             logger.warning(f"Set device({device}) to 'cpu' for only cpu available.")
             self.device = "cpu"
         self.device = paddle.device.set_device(self.device)

ppsci/utils/callbacks.py

@@ -99,7 +99,8 @@ def on_job_start(self, config: DictConfig, **kwargs: Any) -> None:
         if "device" in full_cfg:
             import paddle
 
-            paddle.device.set_device(full_cfg.device)
+            if isinstance(full_cfg.device, str):
+                paddle.device.set_device(full_cfg.device)
 
         # enable prim if specified
         if "prim" in full_cfg and bool(full_cfg.prim):

ppsci/utils/config.py

@@ -305,7 +305,7 @@ class SolverConfig(BaseModel):
             use_tbd: bool = False
             wandb_config: Mapping = {}
             use_wandb: bool = False
-            device: Literal["cpu", "gpu", "xpu"] = "gpu"
+            device: Literal["cpu", "gpu", "xpu", None] = None
             use_amp: bool = False
             amp_level: Literal["O0", "O1", "O2", "OD"] = "O1"
             to_static: bool = False