PyTorch:torch.cuda

该包增加了对CUDA张量类型的支持,实现了与CPU张量相同的功能,但使用GPU进行计算。它是懒惰的初始化,所以你可以随时导入它,并使用is_available()来确定系统是否支持CUDA。CUDA语义中有关于使用CUDA的更多细节。返回cublasHandle_t指针,指向当前cuBLAS句柄返回当前所选设备的索引。返回一个当前所选的Stream上下文管理器,可以更改所选设备。参数:idx (i

PyTorch:CUDA语义

torch.cuda会记录当前选择的GPU,并且分配的所有CUDA张量将在上面创建。可以使用torch.cuda.device上下文管理器更改所选设备。但是,一旦张量被分配,您可以直接对其进行操作,而不考虑所选择的设备,结果将始终放在与张量相同的设备上。默认情况下,不支持跨GPU操作,唯一的例外是copy_()。 除非启用对等存储器访问,否则对分布不同设备上的张量任何启动操作的尝试都将会引发错误。