PyTorch torch.set_default_device 函数
Pytorch torch 参考手册torch.set_default_device 是 PyTorch 中用于设置默认设备的函数。
函数定义
torch.set_default_device(device)
使用示例
实例
import torch
# 设置默认设备
if torch.cuda.is_available():
torch.set_default_device('cuda')
x = torch.zeros(3, 4)
print("默认设备:", x.device)
torch.set_default_device('cpu')
else:
print("CUDA 不可用")
# 设置默认设备
if torch.cuda.is_available():
torch.set_default_device('cuda')
x = torch.zeros(3, 4)
print("默认设备:", x.device)
torch.set_default_device('cpu')
else:
print("CUDA 不可用")