torch.cuda用法

下面是torch.cuda的一些常用函数:

import torch
#检查cuda是否可以使用
torch.cuda.is_available()
#查看当前gpu索引号
torch.cuda.current_device()
#查看当前cuda流
torch.cuda.current_stream(device=0)
#选择device
torch.cuda.device(1)
#查看有多少个GPU设备
torch.cuda.device_count()
#查看gpu的容量
torch.cuda.get_device_capability(device=0)

你可能感兴趣的:(cuda,pytorch,gpu,深度学习)