Unverified Commit e6309acd authored by Jane (Yuan) Xu's avatar Jane (Yuan) Xu Committed by GitHub
Browse files

Simplify `from_blob` usage in `get_cuda_view_from_cpu_tensor` (#29027)


Signed-off-by: default avatarJane Xu <janeyx@meta.com>
parent 988ee66b
...@@ -22,15 +22,10 @@ torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor) { ...@@ -22,15 +22,10 @@ torch::Tensor get_cuda_view_from_cpu_tensor(torch::Tensor& cpu_tensor) {
auto strides = cpu_tensor.strides(); auto strides = cpu_tensor.strides();
auto options = cpu_tensor.options().device(torch::kCUDA); auto options = cpu_tensor.options().device(torch::kCUDA);
// from_blob signature: from_blob(void *data, IntArrayRef sizes, ..., Deleter, // use default no-op deleter, since the memory is owned by the original CPU
// const TensorOptions &) Provide a no-op deleter. The CPU tensor holds the // tensor
// memory, so we don't free it here.
auto deleter = [](void*) {
// no-op, since the memory is owned by the original CPU tensor
};
torch::Tensor cuda_tensor = torch::Tensor cuda_tensor =
torch::from_blob(device_ptr, sizes, strides, deleter, options); torch::from_blob(device_ptr, sizes, strides, options);
TORCH_CHECK(cuda_tensor.device().is_cuda(), TORCH_CHECK(cuda_tensor.device().is_cuda(),
"Resulting tensor is not on CUDA device"); "Resulting tensor is not on CUDA device");
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment