Add at method to TensorListGPU (NVIDIA#131)

* Added at method for TensorListGPU Signed-off-by: ptredak <ptredak@nvidia.com> * Fix Signed-off-by: ptredak <ptredak@nvidia.com> Signed-off-by: Andrei <andreii@nvidia.com>
drivanov · Sep 4, 2018 · 8d5b308 · 8d5b308
1 parent 645cfb0
commit 8d5b308
Showing 1 changed file with 13 additions and 0 deletions.
diff --git a/dali/python/backend_impl.cc b/dali/python/backend_impl.cc
@@ -352,6 +352,19 @@ void ExposeTensorList(py::module &m) { // NOLINT
       Parameters
       ----------
       )code")
+    .def("at",
+        [](TensorList<GPUBackend> &t, Index id) -> std::unique_ptr<Tensor<GPUBackend>> {
+          std::unique_ptr<Tensor<GPUBackend>> ptr(new Tensor<GPUBackend>());
+          ptr->ShareData(&t, id);
+          return ptr;
+        },
+      R"code(
+      Returns a tensor at given position in the list.
+
+      Parameters
+      ----------
+      )code",
+      py::keep_alive<0,1>())
     .def("as_tensor", &TensorList<GPUBackend>::AsTensor,
       R"code(
       Returns a tensor that is a view of this `TensorList`.