Add a currently unexported torch_tensor_free that deletes tensors w…

…ithout requiring us to wait for GC. (#1194)
mlverse · Sep 20, 2024 · 030cda0 · 030cda0
1 parent 7142cf7
commit 030cda0
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git a/src/tensor.cpp b/src/tensor.cpp
@@ -416,4 +416,9 @@ torch::Tensor nnf_pad_circular(torch::Tensor input,
 // [[Rcpp::export]]
 bool cpp_method_Tensor_is_sparse (torch::Tensor x) {
   return lantern_Tensor_is_sparse(x.get());
+}
+
+// [[Rcpp::export]]
+void torch_tensor_free(Rcpp::XPtr<torch::Tensor> x) {
+  x.release();
 }
diff --git a/src/torch_api.cpp b/src/torch_api.cpp
@@ -7,7 +7,10 @@
 
 static inline void tensor_finalizer(SEXP ptr) {
   auto xptr = Rcpp::as<Rcpp::XPtr<XPtrTorchTensor>>(ptr);
-  lantern_tensor_set_pyobj(xptr->get(), nullptr);
+  // When the the xptr is released, we might call this function with an invalid ptr.
+  if (xptr.get()) {
+    lantern_tensor_set_pyobj(xptr->get(), nullptr);
+  }
 }
 
 static inline bool rlang_is_named(SEXP x) {
@@ -140,7 +143,9 @@ XPtrTorchTensor from_sexp_tensor(SEXP x) {
   Rcpp::stop("Expected a torch_tensor.");
 }
 
-void delete_tensor(void* x) { lantern_Tensor_delete(x); }
+void delete_tensor(void* x) { 
+  lantern_Tensor_delete(x); 
+}
 
 // optional_torch_tensor