Browse Source

perf(tensor): always copy to pinned memory before h2d

GitOrigin-RevId: f945f2435f
release-1.4
Megvii Engine Team 4 years ago
parent
commit
3bfe87ecb5
1 changed files with 1 additions and 4 deletions
  1. +1
    -4
      imperative/python/src/tensor.cpp

+ 1
- 4
imperative/python/src/tensor.cpp View File

@@ -271,10 +271,7 @@ TensorWrapper::TensorWrapper(PyObject* args, PyObject* kwargs) {
}

interpreter::Interpreter::Handle handle;
constexpr auto size_threshhold = TensorShape::MAX_NDIM;
if (data.size() > size_threshhold) {
handle = interpreter_for_py->put(npy::np2tensor(data.ptr(), npy::Meth::borrow(cn), dtype), no_cache);
} else {
{
HostTensorND ret(cn);
handle = interpreter_for_py->put(npy::np2tensor(data.ptr(), npy::Meth::copy_into(&ret), dtype), no_cache);
}


Loading…
Cancel
Save