Merge pull request #5904 from longjon/gpu-ptr

Expose GPU pointers to Python
chuckcho · Sep 19, 2017 · effcdb0 · effcdb0
2 parents ea455eb + 8bdc87f
commit effcdb0
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/python/caffe/_caffe.cpp b/python/caffe/_caffe.cpp
@@ -464,6 +464,14 @@ BOOST_PYTHON_MODULE(_caffe) {
  .add_property("count", static_cast<int (Blob<Dtype>::*)() const>(
  &Blob<Dtype>::count))
  .def("reshape", bp::raw_function(&Blob_Reshape))
+#ifndef CPU_ONLY
+ .add_property("_gpu_data_ptr",
+ reinterpret_cast<uintptr_t (Blob<Dtype>::*)()>(
+ &Blob<Dtype>::mutable_gpu_data))
+ .add_property("_gpu_diff_ptr",
+ reinterpret_cast<uintptr_t (Blob<Dtype>::*)()>(
+ &Blob<Dtype>::mutable_gpu_diff))
+#endif
  .add_property("data", bp::make_function(&Blob<Dtype>::mutable_cpu_data,
  NdarrayCallPolicies()))
  .add_property("diff", bp::make_function(&Blob<Dtype>::mutable_cpu_diff,