diff --git a/csrc/cpu/metis_cpu.cpp b/csrc/cpu/metis_cpu.cpp
index 8de38b51..51933f68 100644
--- a/csrc/cpu/metis_cpu.cpp
+++ b/csrc/cpu/metis_cpu.cpp
@@ -11,8 +11,8 @@
 #include "utils.h"
 
 torch::Tensor partition_cpu(torch::Tensor rowptr, torch::Tensor col,
-                            torch::optional<torch::Tensor> optional_value,
-                            torch::optional<torch::Tensor> optional_node_weight,
+                            std::optional<torch::Tensor> optional_value,
+                            std::optional<torch::Tensor> optional_node_weight,
                             int64_t num_parts, bool recursive) {
 #ifdef WITH_METIS
   CHECK_CPU(rowptr);
@@ -66,8 +66,8 @@ torch::Tensor partition_cpu(torch::Tensor rowptr, torch::Tensor col,
 //             --partitions64bit
 torch::Tensor
 mt_partition_cpu(torch::Tensor rowptr, torch::Tensor col,
-                 torch::optional<torch::Tensor> optional_value,
-                 torch::optional<torch::Tensor> optional_node_weight,
+                 std::optional<torch::Tensor> optional_value,
+                 std::optional<torch::Tensor> optional_node_weight,
                  int64_t num_parts, bool recursive, int64_t num_workers) {
 #ifdef WITH_MTMETIS
   CHECK_CPU(rowptr);
diff --git a/csrc/cpu/metis_cpu.h b/csrc/cpu/metis_cpu.h
index 808c4024..685d30ce 100644
--- a/csrc/cpu/metis_cpu.h
+++ b/csrc/cpu/metis_cpu.h
@@ -3,12 +3,12 @@
 #include "../extensions.h"
 
 torch::Tensor partition_cpu(torch::Tensor rowptr, torch::Tensor col,
-                            torch::optional<torch::Tensor> optional_value,
-                            torch::optional<torch::Tensor> optional_node_weight,
+                            std::optional<torch::Tensor> optional_value,
+                            std::optional<torch::Tensor> optional_node_weight,
                             int64_t num_parts, bool recursive);
 
 torch::Tensor
 mt_partition_cpu(torch::Tensor rowptr, torch::Tensor col,
-                 torch::optional<torch::Tensor> optional_value,
-                 torch::optional<torch::Tensor> optional_node_weight,
+                 std::optional<torch::Tensor> optional_value,
+                 std::optional<torch::Tensor> optional_node_weight,
                  int64_t num_parts, bool recursive, int64_t num_workers);
diff --git a/csrc/cpu/relabel_cpu.cpp b/csrc/cpu/relabel_cpu.cpp
index 6dc94c80..db40bbd6 100644
--- a/csrc/cpu/relabel_cpu.cpp
+++ b/csrc/cpu/relabel_cpu.cpp
@@ -42,10 +42,10 @@ std::tuple<torch::Tensor, torch::Tensor> relabel_cpu(torch::Tensor col,
   return std::make_tuple(out_col, out_idx);
 }
 
-std::tuple<torch::Tensor, torch::Tensor, torch::optional<torch::Tensor>,
+std::tuple<torch::Tensor, torch::Tensor, std::optional<torch::Tensor>,
            torch::Tensor>
 relabel_one_hop_cpu(torch::Tensor rowptr, torch::Tensor col,
-                    torch::optional<torch::Tensor> optional_value,
+                    std::optional<torch::Tensor> optional_value,
                     torch::Tensor idx, bool bipartite) {
 
   CHECK_CPU(rowptr);
@@ -79,7 +79,7 @@ relabel_one_hop_cpu(torch::Tensor rowptr, torch::Tensor col,
   auto out_col = torch::empty({offset}, col.options());
   auto out_col_data = out_col.data_ptr<int64_t>();
 
-  torch::optional<torch::Tensor> out_value = torch::nullopt;
+  std::optional<torch::Tensor> out_value = std::nullopt;
   if (optional_value.has_value()) {
     out_value = torch::empty({offset}, optional_value.value().options());
 
diff --git a/csrc/cpu/relabel_cpu.h b/csrc/cpu/relabel_cpu.h
index 999f910a..f10ae043 100644
--- a/csrc/cpu/relabel_cpu.h
+++ b/csrc/cpu/relabel_cpu.h
@@ -5,8 +5,8 @@
 std::tuple<torch::Tensor, torch::Tensor> relabel_cpu(torch::Tensor col,
                                                      torch::Tensor idx);
 
-std::tuple<torch::Tensor, torch::Tensor, torch::optional<torch::Tensor>,
+std::tuple<torch::Tensor, torch::Tensor, std::optional<torch::Tensor>,
            torch::Tensor>
 relabel_one_hop_cpu(torch::Tensor rowptr, torch::Tensor col,
-                    torch::optional<torch::Tensor> optional_value,
+                    std::optional<torch::Tensor> optional_value,
                     torch::Tensor idx, bool bipartite);
diff --git a/csrc/cpu/spmm_cpu.cpp b/csrc/cpu/spmm_cpu.cpp
index a7f2be61..ff54c44f 100644
--- a/csrc/cpu/spmm_cpu.cpp
+++ b/csrc/cpu/spmm_cpu.cpp
@@ -5,9 +5,9 @@
 #include "reducer.h"
 #include "utils.h"
 
-std::tuple<torch::Tensor, torch::optional<torch::Tensor>>
+std::tuple<torch::Tensor, std::optional<torch::Tensor>>
 spmm_cpu(torch::Tensor rowptr, torch::Tensor col,
-         torch::optional<torch::Tensor> optional_value, torch::Tensor mat,
+         std::optional<torch::Tensor> optional_value, torch::Tensor mat,
          std::string reduce) {
   CHECK_CPU(rowptr);
   CHECK_CPU(col);
@@ -29,7 +29,7 @@ spmm_cpu(torch::Tensor rowptr, torch::Tensor col,
   sizes[mat.dim() - 2] = rowptr.numel() - 1;
   auto out = torch::empty(sizes, mat.options());
 
-  torch::optional<torch::Tensor> arg_out = torch::nullopt;
+  std::optional<torch::Tensor> arg_out = std::nullopt;
   int64_t *arg_out_data = nullptr;
   if (reduce2REDUCE.at(reduce) == MIN || reduce2REDUCE.at(reduce) == MAX) {
     arg_out = torch::full_like(out, col.numel(), rowptr.options());
diff --git a/csrc/cpu/spmm_cpu.h b/csrc/cpu/spmm_cpu.h
index b8953844..c709437c 100644
--- a/csrc/cpu/spmm_cpu.h
+++ b/csrc/cpu/spmm_cpu.h
@@ -2,9 +2,9 @@
 
 #include "../extensions.h"
 
-std::tuple<torch::Tensor, torch::optional<torch::Tensor>>
+std::tuple<torch::Tensor, std::optional<torch::Tensor>>
 spmm_cpu(torch::Tensor rowptr, torch::Tensor col,
-         torch::optional<torch::Tensor> optional_value, torch::Tensor mat,
+         std::optional<torch::Tensor> optional_value, torch::Tensor mat,
          std::string reduce);
 
 torch::Tensor spmm_value_bw_cpu(torch::Tensor row, torch::Tensor rowptr,
diff --git a/csrc/cpu/utils.h b/csrc/cpu/utils.h
index b3866948..3c7eebb8 100644
--- a/csrc/cpu/utils.h
+++ b/csrc/cpu/utils.h
@@ -51,7 +51,7 @@ inline int64_t uniform_randint(int64_t high) {
 
 inline torch::Tensor
 choice(int64_t population, int64_t num_samples, bool replace = false,
-       torch::optional<torch::Tensor> weight = torch::nullopt) {
+       std::optional<torch::Tensor> weight = std::nullopt) {
 
   if (population == 0 || num_samples == 0)
     return torch::empty({0}, at::kLong);
diff --git a/csrc/cuda/spmm_cuda.cu b/csrc/cuda/spmm_cuda.cu
index f67b560c..6f48afae 100644
--- a/csrc/cuda/spmm_cuda.cu
+++ b/csrc/cuda/spmm_cuda.cu
@@ -89,9 +89,9 @@ __global__ void spmm_kernel(const int64_t *rowptr_data, const int64_t *col_data,
   }
 }
 
-std::tuple<torch::Tensor, torch::optional<torch::Tensor>>
+std::tuple<torch::Tensor, std::optional<torch::Tensor>>
 spmm_cuda(torch::Tensor rowptr, torch::Tensor col,
-          torch::optional<torch::Tensor> optional_value, torch::Tensor mat,
+          std::optional<torch::Tensor> optional_value, torch::Tensor mat,
           std::string reduce) {
 
   CHECK_CUDA(rowptr);
@@ -115,7 +115,7 @@ spmm_cuda(torch::Tensor rowptr, torch::Tensor col,
   sizes[mat.dim() - 2] = rowptr.numel() - 1;
   auto out = torch::empty(sizes, mat.options());
 
-  torch::optional<torch::Tensor> arg_out = torch::nullopt;
+  std::optional<torch::Tensor> arg_out = std::nullopt;
   int64_t *arg_out_data = nullptr;
   if (reduce2REDUCE.at(reduce) == MIN || reduce2REDUCE.at(reduce) == MAX) {
     arg_out = torch::full_like(out, col.numel(), rowptr.options());
diff --git a/csrc/cuda/spmm_cuda.h b/csrc/cuda/spmm_cuda.h
index 09a11685..421f427d 100644
--- a/csrc/cuda/spmm_cuda.h
+++ b/csrc/cuda/spmm_cuda.h
@@ -2,9 +2,9 @@
 
 #include "../extensions.h"
 
-std::tuple<torch::Tensor, torch::optional<torch::Tensor>>
+std::tuple<torch::Tensor, std::optional<torch::Tensor>>
 spmm_cuda(torch::Tensor rowptr, torch::Tensor col,
-          torch::optional<torch::Tensor> optional_value, torch::Tensor mat,
+          std::optional<torch::Tensor> optional_value, torch::Tensor mat,
           std::string reduce);
 
 torch::Tensor spmm_value_bw_cuda(torch::Tensor row, torch::Tensor rowptr,
diff --git a/csrc/metis.cpp b/csrc/metis.cpp
index 1911e823..084a42f4 100644
--- a/csrc/metis.cpp
+++ b/csrc/metis.cpp
@@ -16,7 +16,7 @@ PyMODINIT_FUNC PyInit__metis_cpu(void) { return NULL; }
 #endif
 
 SPARSE_API torch::Tensor partition(torch::Tensor rowptr, torch::Tensor col,
-                        torch::optional<torch::Tensor> optional_value,
+                        std::optional<torch::Tensor> optional_value,
                         int64_t num_parts, bool recursive) {
   if (rowptr.device().is_cuda()) {
 #ifdef WITH_CUDA
@@ -25,14 +25,14 @@ SPARSE_API torch::Tensor partition(torch::Tensor rowptr, torch::Tensor col,
     AT_ERROR("Not compiled with CUDA support");
 #endif
   } else {
-    return partition_cpu(rowptr, col, optional_value, torch::nullopt, num_parts,
+    return partition_cpu(rowptr, col, optional_value, std::nullopt, num_parts,
                          recursive);
   }
 }
 
 SPARSE_API torch::Tensor partition2(torch::Tensor rowptr, torch::Tensor col,
-                         torch::optional<torch::Tensor> optional_value,
-                         torch::optional<torch::Tensor> optional_node_weight,
+                         std::optional<torch::Tensor> optional_value,
+                         std::optional<torch::Tensor> optional_node_weight,
                          int64_t num_parts, bool recursive) {
   if (rowptr.device().is_cuda()) {
 #ifdef WITH_CUDA
@@ -47,8 +47,8 @@ SPARSE_API torch::Tensor partition2(torch::Tensor rowptr, torch::Tensor col,
 }
 
 SPARSE_API torch::Tensor mt_partition(torch::Tensor rowptr, torch::Tensor col,
-                           torch::optional<torch::Tensor> optional_value,
-                           torch::optional<torch::Tensor> optional_node_weight,
+                           std::optional<torch::Tensor> optional_value,
+                           std::optional<torch::Tensor> optional_node_weight,
                            int64_t num_parts, bool recursive,
                            int64_t num_workers) {
   if (rowptr.device().is_cuda()) {
diff --git a/csrc/relabel.cpp b/csrc/relabel.cpp
index 9f2d4909..aeb20217 100644
--- a/csrc/relabel.cpp
+++ b/csrc/relabel.cpp
@@ -28,10 +28,10 @@ SPARSE_API std::tuple<torch::Tensor, torch::Tensor> relabel(torch::Tensor col,
   }
 }
 
-SPARSE_API std::tuple<torch::Tensor, torch::Tensor, torch::optional<torch::Tensor>,
+SPARSE_API std::tuple<torch::Tensor, torch::Tensor, std::optional<torch::Tensor>,
            torch::Tensor>
 relabel_one_hop(torch::Tensor rowptr, torch::Tensor col,
-                torch::optional<torch::Tensor> optional_value,
+                std::optional<torch::Tensor> optional_value,
                 torch::Tensor idx, bool bipartite) {
   if (rowptr.device().is_cuda()) {
 #ifdef WITH_CUDA
diff --git a/csrc/sparse.h b/csrc/sparse.h
index 10272431..56a6630d 100644
--- a/csrc/sparse.h
+++ b/csrc/sparse.h
@@ -16,28 +16,28 @@ SPARSE_API torch::Tensor ptr2ind(torch::Tensor ptr, int64_t E);
 
 SPARSE_API torch::Tensor
 partition(torch::Tensor rowptr, torch::Tensor col,
-          torch::optional<torch::Tensor> optional_value, int64_t num_parts,
+          std::optional<torch::Tensor> optional_value, int64_t num_parts,
           bool recursive);
 
 SPARSE_API torch::Tensor
 partition2(torch::Tensor rowptr, torch::Tensor col,
-           torch::optional<torch::Tensor> optional_value,
-           torch::optional<torch::Tensor> optional_node_weight,
+           std::optional<torch::Tensor> optional_value,
+           std::optional<torch::Tensor> optional_node_weight,
            int64_t num_parts, bool recursive);
 
 SPARSE_API torch::Tensor
 mt_partition(torch::Tensor rowptr, torch::Tensor col,
-             torch::optional<torch::Tensor> optional_value,
-             torch::optional<torch::Tensor> optional_node_weight,
+             std::optional<torch::Tensor> optional_value,
+             std::optional<torch::Tensor> optional_node_weight,
              int64_t num_parts, bool recursive, int64_t num_workers);
 
 SPARSE_API std::tuple<torch::Tensor, torch::Tensor> relabel(torch::Tensor col,
                                                             torch::Tensor idx);
 
 SPARSE_API std::tuple<torch::Tensor, torch::Tensor,
-                      torch::optional<torch::Tensor>, torch::Tensor>
+                      std::optional<torch::Tensor>, torch::Tensor>
 relabel_one_hop(torch::Tensor rowptr, torch::Tensor col,
-                torch::optional<torch::Tensor> optional_value,
+                std::optional<torch::Tensor> optional_value,
                 torch::Tensor idx, bool bipartite);
 
 SPARSE_API torch::Tensor random_walk(torch::Tensor rowptr, torch::Tensor col,
@@ -52,25 +52,25 @@ std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
 sample_adj(torch::Tensor rowptr, torch::Tensor col, torch::Tensor idx,
            int64_t num_neighbors, bool replace);
 
-SPARSE_API torch::Tensor spmm_sum(torch::optional<torch::Tensor> opt_row,
+SPARSE_API torch::Tensor spmm_sum(std::optional<torch::Tensor> opt_row,
                                   torch::Tensor rowptr, torch::Tensor col,
-                                  torch::optional<torch::Tensor> opt_value,
-                                  torch::optional<torch::Tensor> opt_colptr,
-                                  torch::optional<torch::Tensor> opt_csr2csc,
+                                  std::optional<torch::Tensor> opt_value,
+                                  std::optional<torch::Tensor> opt_colptr,
+                                  std::optional<torch::Tensor> opt_csr2csc,
                                   torch::Tensor mat);
 
-SPARSE_API torch::Tensor spmm_mean(torch::optional<torch::Tensor> opt_row,
+SPARSE_API torch::Tensor spmm_mean(std::optional<torch::Tensor> opt_row,
                                    torch::Tensor rowptr, torch::Tensor col,
-                                   torch::optional<torch::Tensor> opt_value,
-                                   torch::optional<torch::Tensor> opt_rowcount,
-                                   torch::optional<torch::Tensor> opt_colptr,
-                                   torch::optional<torch::Tensor> opt_csr2csc,
+                                   std::optional<torch::Tensor> opt_value,
+                                   std::optional<torch::Tensor> opt_rowcount,
+                                   std::optional<torch::Tensor> opt_colptr,
+                                   std::optional<torch::Tensor> opt_csr2csc,
                                    torch::Tensor mat);
 
 SPARSE_API std::tuple<torch::Tensor, torch::Tensor>
 spmm_min(torch::Tensor rowptr, torch::Tensor col,
-         torch::optional<torch::Tensor> opt_value, torch::Tensor mat);
+         std::optional<torch::Tensor> opt_value, torch::Tensor mat);
 
 SPARSE_API std::tuple<torch::Tensor, torch::Tensor>
 spmm_max(torch::Tensor rowptr, torch::Tensor col,
-         torch::optional<torch::Tensor> opt_value, torch::Tensor mat);
+         std::optional<torch::Tensor> opt_value, torch::Tensor mat);
diff --git a/csrc/spmm.cpp b/csrc/spmm.cpp
index 66fa98cc..ee6f56ed 100644
--- a/csrc/spmm.cpp
+++ b/csrc/spmm.cpp
@@ -19,9 +19,9 @@ PyMODINIT_FUNC PyInit__spmm_cpu(void) { return NULL; }
 #endif
 #endif
 
-std::tuple<torch::Tensor, torch::optional<torch::Tensor>>
+std::tuple<torch::Tensor, std::optional<torch::Tensor>>
 spmm_fw(torch::Tensor rowptr, torch::Tensor col,
-        torch::optional<torch::Tensor> optional_value, torch::Tensor mat,
+        std::optional<torch::Tensor> optional_value, torch::Tensor mat,
         std::string reduce) {
   if (rowptr.device().is_cuda()) {
 #ifdef WITH_CUDA
@@ -55,10 +55,10 @@ using torch::autograd::variable_list;
 class SPMMSum : public torch::autograd::Function<SPMMSum> {
 public:
   static variable_list forward(AutogradContext *ctx,
-                               torch::optional<Variable> opt_row,
+                               std::optional<Variable> opt_row,
                                Variable rowptr, Variable col, Variable value,
-                               torch::optional<Variable> opt_colptr,
-                               torch::optional<Variable> opt_csr2csc,
+                               std::optional<Variable> opt_colptr,
+                               std::optional<Variable> opt_csr2csc,
                                Variable mat, bool has_value) {
 
     if (has_value && torch::autograd::any_variable_requires_grad({value})) {
@@ -75,7 +75,7 @@ class SPMMSum : public torch::autograd::Function<SPMMSum> {
     auto colptr = opt_colptr.has_value() ? opt_colptr.value() : col;
     auto csr2csc = opt_csr2csc.has_value() ? opt_csr2csc.value() : col;
 
-    torch::optional<torch::Tensor> opt_value = torch::nullopt;
+    std::optional<torch::Tensor> opt_value = std::nullopt;
     if (has_value)
       opt_value = value;
 
@@ -99,7 +99,7 @@ class SPMMSum : public torch::autograd::Function<SPMMSum> {
 
     auto grad_mat = Variable();
     if (torch::autograd::any_variable_requires_grad({mat})) {
-      torch::optional<torch::Tensor> opt_value = torch::nullopt;
+      std::optional<torch::Tensor> opt_value = std::nullopt;
       if (has_value)
         opt_value = value.view({-1, 1}).index_select(0, csr2csc).view(-1);
 
@@ -115,11 +115,11 @@ class SPMMSum : public torch::autograd::Function<SPMMSum> {
 class SPMMMean : public torch::autograd::Function<SPMMMean> {
 public:
   static variable_list forward(AutogradContext *ctx,
-                               torch::optional<Variable> opt_row,
+                               std::optional<Variable> opt_row,
                                Variable rowptr, Variable col, Variable value,
-                               torch::optional<Variable> opt_rowcount,
-                               torch::optional<Variable> opt_colptr,
-                               torch::optional<Variable> opt_csr2csc,
+                               std::optional<Variable> opt_rowcount,
+                               std::optional<Variable> opt_colptr,
+                               std::optional<Variable> opt_csr2csc,
                                Variable mat, bool has_value) {
 
     if (has_value && torch::autograd::any_variable_requires_grad({value})) {
@@ -138,7 +138,7 @@ class SPMMMean : public torch::autograd::Function<SPMMMean> {
     auto colptr = opt_colptr.has_value() ? opt_colptr.value() : col;
     auto csr2csc = opt_csr2csc.has_value() ? opt_csr2csc.value() : col;
 
-    torch::optional<torch::Tensor> opt_value = torch::nullopt;
+    std::optional<torch::Tensor> opt_value = std::nullopt;
     if (has_value)
       opt_value = value;
 
@@ -188,7 +188,7 @@ class SPMMMin : public torch::autograd::Function<SPMMMin> {
                                Variable col, Variable value, Variable mat,
                                bool has_value) {
 
-    torch::optional<torch::Tensor> opt_value = torch::nullopt;
+    std::optional<torch::Tensor> opt_value = std::nullopt;
     if (has_value)
       opt_value = value;
 
@@ -248,7 +248,7 @@ class SPMMMax : public torch::autograd::Function<SPMMMax> {
                                Variable col, Variable value, Variable mat,
                                bool has_value) {
 
-    torch::optional<torch::Tensor> opt_value = torch::nullopt;
+    std::optional<torch::Tensor> opt_value = std::nullopt;
     if (has_value)
       opt_value = value;
 
@@ -302,23 +302,23 @@ class SPMMMax : public torch::autograd::Function<SPMMMax> {
   }
 };
 
-SPARSE_API torch::Tensor spmm_sum(torch::optional<torch::Tensor> opt_row,
+SPARSE_API torch::Tensor spmm_sum(std::optional<torch::Tensor> opt_row,
                        torch::Tensor rowptr, torch::Tensor col,
-                       torch::optional<torch::Tensor> opt_value,
-                       torch::optional<torch::Tensor> opt_colptr,
-                       torch::optional<torch::Tensor> opt_csr2csc,
+                       std::optional<torch::Tensor> opt_value,
+                       std::optional<torch::Tensor> opt_colptr,
+                       std::optional<torch::Tensor> opt_csr2csc,
                        torch::Tensor mat) {
   auto value = opt_value.has_value() ? opt_value.value() : col;
   return SPMMSum::apply(opt_row, rowptr, col, value, opt_colptr, opt_csr2csc,
                         mat, opt_value.has_value())[0];
 }
 
-SPARSE_API torch::Tensor spmm_mean(torch::optional<torch::Tensor> opt_row,
+SPARSE_API torch::Tensor spmm_mean(std::optional<torch::Tensor> opt_row,
                         torch::Tensor rowptr, torch::Tensor col,
-                        torch::optional<torch::Tensor> opt_value,
-                        torch::optional<torch::Tensor> opt_rowcount,
-                        torch::optional<torch::Tensor> opt_colptr,
-                        torch::optional<torch::Tensor> opt_csr2csc,
+                        std::optional<torch::Tensor> opt_value,
+                        std::optional<torch::Tensor> opt_rowcount,
+                        std::optional<torch::Tensor> opt_colptr,
+                        std::optional<torch::Tensor> opt_csr2csc,
                         torch::Tensor mat) {
   auto value = opt_value.has_value() ? opt_value.value() : col;
   return SPMMMean::apply(opt_row, rowptr, col, value, opt_rowcount, opt_colptr,
@@ -327,7 +327,7 @@ SPARSE_API torch::Tensor spmm_mean(torch::optional<torch::Tensor> opt_row,
 
 SPARSE_API std::tuple<torch::Tensor, torch::Tensor>
 spmm_min(torch::Tensor rowptr, torch::Tensor col,
-         torch::optional<torch::Tensor> opt_value, torch::Tensor mat) {
+         std::optional<torch::Tensor> opt_value, torch::Tensor mat) {
   auto value = opt_value.has_value() ? opt_value.value() : col;
   auto result = SPMMMin::apply(rowptr, col, value, mat, opt_value.has_value());
   return std::make_tuple(result[0], result[1]);
@@ -335,7 +335,7 @@ spmm_min(torch::Tensor rowptr, torch::Tensor col,
 
 SPARSE_API std::tuple<torch::Tensor, torch::Tensor>
 spmm_max(torch::Tensor rowptr, torch::Tensor col,
-         torch::optional<torch::Tensor> opt_value, torch::Tensor mat) {
+         std::optional<torch::Tensor> opt_value, torch::Tensor mat) {
   auto value = opt_value.has_value() ? opt_value.value() : col;
   auto result = SPMMMax::apply(rowptr, col, value, mat, opt_value.has_value());
   return std::make_tuple(result[0], result[1]);