From e5e1316aa790cf8bd8cdff8316cba9ac409219a4 Mon Sep 17 00:00:00 2001
From: Muhammed Fatih Balin <m.f.balin@gmail.com>
Date: Sat, 6 Jul 2024 23:29:39 +0000
Subject: [PATCH] result tensor options should come from input.

---
 graphbolt/src/cnumpy.cc | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/graphbolt/src/cnumpy.cc b/graphbolt/src/cnumpy.cc
index b4e04be7d3a6..0e68893f474c 100644
--- a/graphbolt/src/cnumpy.cc
+++ b/graphbolt/src/cnumpy.cc
@@ -127,10 +127,9 @@ torch::Tensor OnDiskNpyArray::IndexSelectIOUring(torch::Tensor index) {
   shape.push_back(num_index);
   shape.insert(shape.end(), feature_dim_.begin() + 1, feature_dim_.end());
   auto result = torch::empty(
-      shape, torch::TensorOptions()
+      shape, index.options()
                  .dtype(dtype_)
                  .layout(torch::kStrided)
-                 .device(torch::kCPU)
                  .requires_grad(false));
   auto result_buffer = reinterpret_cast<char *>(result.data_ptr());