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 記憶體管理的更多詳細資訊，請參閱 記憶體管理。