We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
1 parent cc85284 commit 218cdc9Copy full SHA for 218cdc9
1 file changed
deepmd/utils/summary.py
@@ -79,16 +79,19 @@ def __call__(self) -> None:
79
try:
80
if backend == "PyTorch":
81
import torch
82
+
83
if torch.cuda.is_available():
84
device_name = torch.cuda.get_device_name(0)
85
elif backend == "TensorFlow":
86
import tensorflow as tf
87
88
gpus = tf.config.list_physical_devices("GPU")
89
if gpus:
90
# Use the first physical GPU device identifier as the device name
91
device_name = gpus[0].name
92
elif backend == "Paddle":
93
import paddle
94
95
# Use Paddle's current device string (e.g., "gpu:0") as a device identifier
96
device_name = paddle.get_device()
97
except Exception:
0 commit comments