Allow access to Data() pointer from device

NVIDIA · Aug 23, 2024 · ea8296a · ea8296a
1 parent a8d8078
commit ea8296a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/include/matx/core/tensor.h b/include/matx/core/tensor.h
@@ -946,7 +946,7 @@ class tensor_t : public detail::tensor_impl_t<T,RANK,Desc> {
    * @returns Underlying data pointer of type T
    *
    */
-  __MATX_HOST__ __MATX_INLINE__ T *Data() const noexcept { return this->ldata_; }
+  __MATX_HOST__ __MATX_DEVICE__ __MATX_INLINE__ T *Data() const noexcept { return this->ldata_; }
 
   /**
    * Set the underlying data pointer from the view
@@ -1473,7 +1473,7 @@ class tensor_t : public detail::tensor_impl_t<T,RANK,Desc> {
       }
     }
 
-    MATX_ASSERT_STR(((RANK - end_count) == N), matxInvalidSize, 
+    MATX_ASSERT_STR(((RANK - end_count) == N), matxInvalidSize,
             "Number of matxDropDim specifiers must match the output rank");
 
 #pragma unroll