From f382849289ab724adb3561799bfa6ab30389a84d Mon Sep 17 00:00:00 2001
From: Sameer Sheorey <sameer.sheorey@intel.com>
Date: Thu, 28 Dec 2023 21:52:39 -0800
Subject: [PATCH] Use correct devices for attr tensors in SelectByIndex.

---
 CMakeLists.txt                         |  6 +++++-
 cpp/open3d/t/geometry/TriangleMesh.cpp | 21 +++++++++------------
 2 files changed, 14 insertions(+), 13 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index c3d64e4b95d0..b571248db475 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -48,7 +48,11 @@ option(BUILD_BENCHMARKS           "Build the micro benchmarks"               OFF
 option(BUILD_PYTHON_MODULE        "Build the python module"                  ON )
 option(BUILD_CUDA_MODULE          "Build the CUDA module"                    OFF)
 option(BUILD_COMMON_CUDA_ARCHS    "Build for common CUDA GPUs (for release)" OFF)
-option(ENABLE_CACHED_CUDA_MANAGER "Enable cached CUDA memory manager"        ON )
+if (WIN32)   # Causes CUDA runtime error on Windows (See issue #6555)
+    option(ENABLE_CACHED_CUDA_MANAGER "Enable cached CUDA memory manager"    OFF)
+else()
+    option(ENABLE_CACHED_CUDA_MANAGER "Enable cached CUDA memory manager"    ON )
+endif()
 if(NOT LINUX_AARCH64 AND NOT APPLE_AARCH64)
     option(BUILD_ISPC_MODULE      "Build the ISPC module"                    ON )
 else()
diff --git a/cpp/open3d/t/geometry/TriangleMesh.cpp b/cpp/open3d/t/geometry/TriangleMesh.cpp
index 907873d1520b..34100c643192 100644
--- a/cpp/open3d/t/geometry/TriangleMesh.cpp
+++ b/cpp/open3d/t/geometry/TriangleMesh.cpp
@@ -1029,8 +1029,8 @@ int TriangleMesh::PCAPartition(int max_faces) {
 }
 
 /// A helper to compute new vertex indices out of vertex mask.
-/// \param tris_cpu tensor with triangle indices to update.
-/// \param vertex_mask tensor with the mask for vertices.
+/// \param tris_cpu CPU tensor with triangle indices to update.
+/// \param vertex_mask CPU tensor with the mask for vertices.
 template <typename T>
 static void UpdateTriangleIndicesByVertexMask(core::Tensor &tris_cpu,
                                               const core::Tensor &vertex_mask) {
@@ -1148,11 +1148,10 @@ static bool IsNegative(T val) {
 }
 
 TriangleMesh TriangleMesh::SelectByIndex(const core::Tensor &indices) const {
-    TriangleMesh result;
     core::AssertTensorShape(indices, {indices.GetLength()});
     if (!HasVertexPositions()) {
         utility::LogWarning("[SelectByIndex] TriangleMesh has no vertices.");
-        return result;
+        return {};
     }
     GetVertexAttr().AssertSizeSynchronized();
 
@@ -1194,7 +1193,7 @@ TriangleMesh TriangleMesh::SelectByIndex(const core::Tensor &indices) const {
                     scalar_tris_t *vertex_mask_ptr =
                             vertex_mask.GetDataPtr<scalar_tris_t>();
                     const scalar_indices_t *indices_ptr =
-                            indices.GetDataPtr<scalar_indices_t>();
+                            indices_cpu.GetDataPtr<scalar_indices_t>();
                     for (int64_t i = 0; i < indices.GetLength(); ++i) {
                         if (IsNegative(indices_ptr[i]) ||
                             indices_ptr[i] >=
@@ -1233,16 +1232,14 @@ TriangleMesh TriangleMesh::SelectByIndex(const core::Tensor &indices) const {
                 });
     });
 
-    // send the vertex mask to original device and apply to vertices
+    // send the vertex mask and triangle mask to original device and apply to
+    // vertices
     vertex_mask = vertex_mask.To(GetDevice(), core::Bool);
+    tri_mask = tri_mask.To(GetDevice());
     core::Tensor new_vertices = GetVertexPositions().IndexGet({vertex_mask});
-    result.SetVertexPositions(new_vertices);
-
-    if (HasTriangleIndices()) {
-        // select triangles and send the selected ones to the original device
-        result.SetTriangleIndices(tris_cpu.To(GetDevice()));
-    }
+    core::Tensor new_tris = tris_cpu.To(GetDevice());
 
+    TriangleMesh result(new_vertices, new_tris);
     CopyAttributesByMasks(result, *this, vertex_mask, tri_mask);
 
     return result;