/usr/local/lib64/python3.6/site-packages/torch/include/THC
Edit: /usr/local/lib64/python3.6/site-packages/torch/include/THC/THCThrustAllocator.cuh (618B)
#ifndef THC_THRUST_ALLOCATOR_INC
#define THC_THRUST_ALLOCATOR_INC
#include
#include
/// Allocator for Thrust to re-route its internal device allocations
/// to the THC allocator
class THCThrustAllocator {
public:
typedef char value_type;
THCThrustAllocator(THCState* state)
: state_(state) {
}
~THCThrustAllocator() {
}
char* allocate(std::ptrdiff_t size) {
return static_cast(THCudaMalloc(state_, size));
}
void deallocate(char* p, size_t size) {
THCudaFree(state_, p);
}
private:
THCState* state_;
};
#endif // THC_THRUST_ALLOCATOR_INC