torch.cuda

该包增加了对CUDA张量类型的支持，实现了与CPU张量相同的功能，但使用GPU进行计算。

它是懒惰的初始化，所以你可以随时导入它，并使用is_available()来确定系统是否支持CUDA。

CUDA语义中有关于使用CUDA的更多细节。

torch.cuda.current_blas_handle()

返回cublasHandle_t指针，指向当前cuBLAS句柄

torch.cuda.current_device()

返回当前所选设备的索引。

torch.cuda.current_stream()

返回一个当前所选的Stream

class torch.cuda.device(idx)

上下文管理器，可以更改所选设备。

参数：

• idx (int) – 设备索引选择。如果这个参数是负的，则是无效操作。

torch.cuda.device_count()

返回可得到的GPU数量。

class torch.cuda.device_of(obj)

将当前设备更改为给定对象的上下文管理器。

可以使用张量和存储作为参数。如果给定的对象不是在GPU上分配的，这是一个无效操作。

参数：

• obj (Tensor or Storage) – 在选定设备上分配的对象。

torch.cuda.is_available()

返回一个bool值，指示CUDA当前是否可用。

torch.cuda.set_device(device)

设置当前设备。

不鼓励使用此函数来设置。在大多数情况下，最好使用CUDA_VISIBLE_DEVICES环境变量。

参数：

• device (int) – 所选设备。如果此参数为负，则此函数是无效操作。

torch.cuda.stream(stream)

选择给定流的上下文管理器。

在其上下文中排队的所有CUDA核心将在所选流上入队。

参数：

• stream (Stream) – 所选流。如果是None，则这个管理器是无效的。

torch.cuda.synchronize()

等待当前设备上所有流中的所有核心完成。