18 lines
519 B
Python
18 lines
519 B
Python
import torch
|
|
|
|
# 检查是否有可用的 GPU
|
|
if torch.cuda.is_available():
|
|
print(f"Number of GPUs available: {torch.cuda.device_count()}")
|
|
for i in range(torch.cuda.device_count()):
|
|
print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
|
|
else:
|
|
print("No GPU available.")
|
|
|
|
# 将张量移动到 GPU 上
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
print(f"Using device: {device}")
|
|
|
|
# 示例：在 GPU 上创建一个张量
|
|
x = torch.tensor([1.0, 2.0, 3.0], device=device)
|
|
print(x)
|