torch.cuda.caching_allocator_alloc

torch.cuda.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 内存管理的详细信息。