PyTorch torch.get_device_module 函数
Pytorch torch 参考手册torch.get_device_module 是 PyTorch 中用于获取当前设备模块的函数。
函数定义
torch.get_device_module(device='cuda')
使用示例
实例
import torch
# 获取CUDA设备模块
cuda_module = torch.get_device_module('cuda')
print("CUDA模块:", cuda_module)
# 检查CUDA是否可用
print("CUDA可用:", torch.cuda.is_available())
# 获取CUDA设备模块
cuda_module = torch.get_device_module('cuda')
print("CUDA模块:", cuda_module)
# 检查CUDA是否可用
print("CUDA可用:", torch.cuda.is_available())