torch.cuda.caching_allocator_alloc

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

使用 CUDA 記憶體分配器執行記憶體分配。

為指定的裝置和流分配記憶體，此函式旨在用於與其他框架的互操作性。分配的記憶體透過 caching_allocator_delete() 釋放。

引數

• size (int) – 要分配的位元組數。

• device (torch.device or int, 可選) – 所選裝置。如果為 None，則使用預設 CUDA 裝置。

• stream (torch.cuda.Stream or int, 可選) – 所選流。如果為 None，則使用所選裝置的預設流。

注意

有關 GPU 記憶體管理的更多詳細資訊，請參閱 記憶體管理。