The CPU memory allocator is implemented here:

  
      github.com
  
  
    pytorch/pytorch/blob/e26fccc22b1f1fa28bceb076567d3e53c92a4c77/c10/core/CPUAllocator.cpp#L100-L120 7

struct C10_API DefaultCPUAllocator final : at::Allocator {  DefaultCPUAllocator() {}  ~DefaultCPUAllocator() override {}  at::DataPtr allocate(size_t nbytes) const override {    void* data = alloc_cpu(nbytes);    profiledCPUMemoryReporter().New(data, nbytes);    return {data, data, &ReportAndDelete, at::Device(at::DeviceType::CPU)};  }  static void ReportAndDelete(void* ptr) {    if (!ptr) {      return;    }    profiledCPUMemoryReporter().Delete(ptr);    free_cpu(ptr);  }  at::DeleterFnPtr raw_deleter() const override {    return &ReportAndDelete;  }};



  
  
    
    
  
  


AFAICT you can control how memory is allocated by implementing another at::Allocator and registering it.