PinnedMemoryAllocator.h 257 B

1234567891011
  1. #pragma once
  2. #include <c10/core/Allocator.h>
  3. #include <ATen/cuda/CachingHostAllocator.h>
  4. namespace at { namespace cuda {
  5. inline TORCH_CUDA_CPP_API at::Allocator* getPinnedMemoryAllocator() {
  6. return getCachingHostAllocator();
  7. }
  8. }} // namespace at::cuda