pytorch_pfn_extras.cuda.use_torch_mempool_in_cupy

pytorch_pfn_extras.cuda.use_torch_mempool_in_cupy()

Use the PyTorch memory pool in CuPy.

If you want to use PyTorch’s memory pool and non-default CUDA streams, streams must be created and managed using PyTorch (using torch.cuda.Stream() and pytorch_pfn_extras.cuda.stream(stream)).

Return type:

None