torch.cuda.memory.caching_allocator_alloc

torch.cuda.memory.caching_allocator_alloc(size, device=None, stream=None)

使用 CUDA 内存分配器执行内存分配。

内存为给定的设备和流进行分配，此函数旨在用于与其他框架进行互操作。分配的内存通过 caching_allocator_delete() 释放。

参数

• size (int) – 要分配的字节数。

• device (torch.device 或 int, 可选) – 选择的设备。如果为 None，则使用默认的 CUDA 设备。

• stream (torch.cuda.Stream 或 int, 可选) – 选择的流。如果为 None，则使用所选设备的默认流。

注意

有关 GPU 内存管理的更多详细信息，请参阅 内存管理。