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, optional) – 選取的裝置。如果為 None，則會使用預設的 CUDA 裝置。

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

注意

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