mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2026-03-16 11:47:09 +08:00
Simplify from_blob usage in get_cuda_view_from_cpu_tensor (#29027)
Signed-off-by: Jane Xu <janeyx@meta.com>
This commit is contained in:
parent
988ee66b0d
commit
e6309acdba
@ -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");
|
||||||
|
|||||||
Loading…
x
Reference in New Issue
Block a user