diff --git a/csrc/cluster.h b/csrc/cluster.h
index b615af6..bd0afdf 100644
--- a/csrc/cluster.h
+++ b/csrc/cluster.h
@@ -14,15 +14,15 @@ CLUSTER_API torch::Tensor fps(torch::Tensor src, torch::Tensor ptr, torch::Tenso
                   bool random_start);
 
 CLUSTER_API torch::Tensor graclus(torch::Tensor rowptr, torch::Tensor col,
-                      torch::optional<torch::Tensor> optional_weight);
+                      std::optional<torch::Tensor> optional_weight);
 
 CLUSTER_API torch::Tensor grid(torch::Tensor pos, torch::Tensor size,
-                   torch::optional<torch::Tensor> optional_start,
-                   torch::optional<torch::Tensor> optional_end);
+                   std::optional<torch::Tensor> optional_start,
+                   std::optional<torch::Tensor> optional_end);
 
 CLUSTER_API torch::Tensor knn(torch::Tensor x, torch::Tensor y,
-                  torch::optional<torch::Tensor> ptr_x,
-                  torch::optional<torch::Tensor> ptr_y, int64_t k, bool cosine,
+                  std::optional<torch::Tensor> ptr_x,
+                  std::optional<torch::Tensor> ptr_y, int64_t k, bool cosine,
                   int64_t num_workers);
 
 CLUSTER_API torch::Tensor nearest(torch::Tensor x, torch::Tensor y, torch::Tensor ptr_x,
diff --git a/csrc/cpu/graclus_cpu.cpp b/csrc/cpu/graclus_cpu.cpp
index e6669d4..f827104 100644
--- a/csrc/cpu/graclus_cpu.cpp
+++ b/csrc/cpu/graclus_cpu.cpp
@@ -3,7 +3,7 @@
 #include "utils.h"
 
 torch::Tensor graclus_cpu(torch::Tensor rowptr, torch::Tensor col,
-                          torch::optional<torch::Tensor> optional_weight) {
+                          std::optional<torch::Tensor> optional_weight) {
   CHECK_CPU(rowptr);
   CHECK_CPU(col);
   CHECK_INPUT(rowptr.dim() == 1 && col.dim() == 1);
diff --git a/csrc/cpu/graclus_cpu.h b/csrc/cpu/graclus_cpu.h
index c3dc83f..3d7442f 100644
--- a/csrc/cpu/graclus_cpu.h
+++ b/csrc/cpu/graclus_cpu.h
@@ -3,4 +3,4 @@
 #include "../extensions.h"
 
 torch::Tensor graclus_cpu(torch::Tensor rowptr, torch::Tensor col,
-                          torch::optional<torch::Tensor> optional_weight);
+                          std::optional<torch::Tensor> optional_weight);
diff --git a/csrc/cpu/grid_cpu.cpp b/csrc/cpu/grid_cpu.cpp
index 7edd708..96f7b45 100644
--- a/csrc/cpu/grid_cpu.cpp
+++ b/csrc/cpu/grid_cpu.cpp
@@ -3,8 +3,8 @@
 #include "utils.h"
 
 torch::Tensor grid_cpu(torch::Tensor pos, torch::Tensor size,
-                       torch::optional<torch::Tensor> optional_start,
-                       torch::optional<torch::Tensor> optional_end) {
+                       std::optional<torch::Tensor> optional_start,
+                       std::optional<torch::Tensor> optional_end) {
 
   CHECK_CPU(pos);
   CHECK_CPU(size);
diff --git a/csrc/cpu/grid_cpu.h b/csrc/cpu/grid_cpu.h
index af7f363..99649d1 100644
--- a/csrc/cpu/grid_cpu.h
+++ b/csrc/cpu/grid_cpu.h
@@ -2,5 +2,5 @@
 
 #include "../extensions.h"
 torch::Tensor grid_cpu(torch::Tensor pos, torch::Tensor size,
-                       torch::optional<torch::Tensor> optional_start,
-                       torch::optional<torch::Tensor> optional_end);
+                       std::optional<torch::Tensor> optional_start,
+                       std::optional<torch::Tensor> optional_end);
diff --git a/csrc/cpu/knn_cpu.cpp b/csrc/cpu/knn_cpu.cpp
index 270f7e3..a67db83 100644
--- a/csrc/cpu/knn_cpu.cpp
+++ b/csrc/cpu/knn_cpu.cpp
@@ -5,8 +5,8 @@
 #include "utils/nanoflann.hpp"
 
 torch::Tensor knn_cpu(torch::Tensor x, torch::Tensor y,
-                      torch::optional<torch::Tensor> ptr_x,
-                      torch::optional<torch::Tensor> ptr_y, int64_t k,
+                      std::optional<torch::Tensor> ptr_x,
+                      std::optional<torch::Tensor> ptr_y, int64_t k,
                       int64_t num_workers) {
 
   CHECK_CPU(x);
diff --git a/csrc/cpu/knn_cpu.h b/csrc/cpu/knn_cpu.h
index 97f11a4..be7d5c0 100644
--- a/csrc/cpu/knn_cpu.h
+++ b/csrc/cpu/knn_cpu.h
@@ -3,6 +3,6 @@
 #include "../extensions.h"
 
 torch::Tensor knn_cpu(torch::Tensor x, torch::Tensor y,
-                      torch::optional<torch::Tensor> ptr_x,
-                      torch::optional<torch::Tensor> ptr_y, int64_t k,
+                      std::optional<torch::Tensor> ptr_x,
+                      std::optional<torch::Tensor> ptr_y, int64_t k,
                       int64_t num_workers);
diff --git a/csrc/cpu/radius_cpu.cpp b/csrc/cpu/radius_cpu.cpp
index 41490f9..3c1b087 100644
--- a/csrc/cpu/radius_cpu.cpp
+++ b/csrc/cpu/radius_cpu.cpp
@@ -5,8 +5,8 @@
 #include "utils/nanoflann.hpp"
 
 torch::Tensor radius_cpu(torch::Tensor x, torch::Tensor y,
-                         torch::optional<torch::Tensor> ptr_x,
-                         torch::optional<torch::Tensor> ptr_y, double r,
+                         std::optional<torch::Tensor> ptr_x,
+                         std::optional<torch::Tensor> ptr_y, double r,
                          int64_t max_num_neighbors, int64_t num_workers,
                          bool ignore_same_index) {
 
diff --git a/csrc/cpu/radius_cpu.h b/csrc/cpu/radius_cpu.h
index 0dbafe5..9b77813 100644
--- a/csrc/cpu/radius_cpu.h
+++ b/csrc/cpu/radius_cpu.h
@@ -3,7 +3,7 @@
 #include "../extensions.h"
 
 torch::Tensor radius_cpu(torch::Tensor x, torch::Tensor y,
-                         torch::optional<torch::Tensor> ptr_x,
-                         torch::optional<torch::Tensor> ptr_y, double r,
+                         std::optional<torch::Tensor> ptr_x,
+                         std::optional<torch::Tensor> ptr_y, double r,
                          int64_t max_num_neighbors, int64_t num_workers,
                          bool ignore_same_index);
diff --git a/csrc/cuda/graclus_cuda.cu b/csrc/cuda/graclus_cuda.cu
index 3bb118b..7743c7a 100644
--- a/csrc/cuda/graclus_cuda.cu
+++ b/csrc/cuda/graclus_cuda.cu
@@ -103,7 +103,7 @@ __global__ void weighted_propose_kernel(int64_t *out, int64_t *proposal,
 
 void propose(torch::Tensor out, torch::Tensor proposal, torch::Tensor rowptr,
              torch::Tensor col,
-             torch::optional<torch::Tensor> optional_weight) {
+             std::optional<torch::Tensor> optional_weight) {
 
   auto stream = at::cuda::getCurrentCUDAStream();
 
@@ -192,7 +192,7 @@ __global__ void weighted_respond_kernel(int64_t *out, const int64_t *proposal,
 
 void respond(torch::Tensor out, torch::Tensor proposal, torch::Tensor rowptr,
              torch::Tensor col,
-             torch::optional<torch::Tensor> optional_weight) {
+             std::optional<torch::Tensor> optional_weight) {
 
   auto stream = at::cuda::getCurrentCUDAStream();
 
@@ -214,7 +214,7 @@ void respond(torch::Tensor out, torch::Tensor proposal, torch::Tensor rowptr,
 }
 
 torch::Tensor graclus_cuda(torch::Tensor rowptr, torch::Tensor col,
-                           torch::optional<torch::Tensor> optional_weight) {
+                           std::optional<torch::Tensor> optional_weight) {
   CHECK_CUDA(rowptr);
   CHECK_CUDA(col);
   CHECK_INPUT(rowptr.dim() == 1 && col.dim() == 1);
diff --git a/csrc/cuda/graclus_cuda.h b/csrc/cuda/graclus_cuda.h
index b9f9358..bfd3fa5 100644
--- a/csrc/cuda/graclus_cuda.h
+++ b/csrc/cuda/graclus_cuda.h
@@ -3,4 +3,4 @@
 #include "../extensions.h"
 
 torch::Tensor graclus_cuda(torch::Tensor rowptr, torch::Tensor col,
-                           torch::optional<torch::Tensor> optional_weight);
+                           std::optional<torch::Tensor> optional_weight);
diff --git a/csrc/cuda/grid_cuda.cu b/csrc/cuda/grid_cuda.cu
index 64037bd..f9e1602 100644
--- a/csrc/cuda/grid_cuda.cu
+++ b/csrc/cuda/grid_cuda.cu
@@ -25,8 +25,8 @@ __global__ void grid_kernel(const scalar_t *pos, const scalar_t *size,
 }
 
 torch::Tensor grid_cuda(torch::Tensor pos, torch::Tensor size,
-                        torch::optional<torch::Tensor> optional_start,
-                        torch::optional<torch::Tensor> optional_end) {
+                        std::optional<torch::Tensor> optional_start,
+                        std::optional<torch::Tensor> optional_end) {
   CHECK_CUDA(pos);
   CHECK_CUDA(size);
   c10::cuda::MaybeSetDevice(pos.get_device());
diff --git a/csrc/cuda/grid_cuda.h b/csrc/cuda/grid_cuda.h
index bfdb984..c63be0c 100644
--- a/csrc/cuda/grid_cuda.h
+++ b/csrc/cuda/grid_cuda.h
@@ -3,5 +3,5 @@
 #include "../extensions.h"
 
 torch::Tensor grid_cuda(torch::Tensor pos, torch::Tensor size,
-                        torch::optional<torch::Tensor> optional_start,
-                        torch::optional<torch::Tensor> optional_end);
+                        std::optional<torch::Tensor> optional_start,
+                        std::optional<torch::Tensor> optional_end);
diff --git a/csrc/cuda/knn_cuda.cu b/csrc/cuda/knn_cuda.cu
index c4dac2a..77acc08 100644
--- a/csrc/cuda/knn_cuda.cu
+++ b/csrc/cuda/knn_cuda.cu
@@ -84,8 +84,8 @@ knn_kernel(const scalar_t *__restrict__ x, const scalar_t *__restrict__ y,
 }
 
 torch::Tensor knn_cuda(const torch::Tensor x, const torch::Tensor y,
-                       torch::optional<torch::Tensor> ptr_x,
-                       torch::optional<torch::Tensor> ptr_y, const int64_t k,
+                       std::optional<torch::Tensor> ptr_x,
+                       std::optional<torch::Tensor> ptr_y, const int64_t k,
                        const bool cosine) {
 
   CHECK_CUDA(x);
diff --git a/csrc/cuda/knn_cuda.h b/csrc/cuda/knn_cuda.h
index 4e732a8..56c1710 100644
--- a/csrc/cuda/knn_cuda.h
+++ b/csrc/cuda/knn_cuda.h
@@ -3,6 +3,6 @@
 #include "../extensions.h"
 
 torch::Tensor knn_cuda(torch::Tensor x, torch::Tensor y,
-                       torch::optional<torch::Tensor> ptr_x,
-                       torch::optional<torch::Tensor> ptr_y, int64_t k,
+                       std::optional<torch::Tensor> ptr_x,
+                       std::optional<torch::Tensor> ptr_y, int64_t k,
                        bool cosine);
diff --git a/csrc/cuda/radius_cuda.cu b/csrc/cuda/radius_cuda.cu
index 29db910..9a18d54 100644
--- a/csrc/cuda/radius_cuda.cu
+++ b/csrc/cuda/radius_cuda.cu
@@ -42,8 +42,8 @@ radius_kernel(const scalar_t *__restrict__ x, const scalar_t *__restrict__ y,
 }
 
 torch::Tensor radius_cuda(const torch::Tensor x, const torch::Tensor y,
-                          torch::optional<torch::Tensor> ptr_x,
-                          torch::optional<torch::Tensor> ptr_y, const double r,
+                          std::optional<torch::Tensor> ptr_x,
+                          std::optional<torch::Tensor> ptr_y, const double r,
                           const int64_t max_num_neighbors,
                           const bool ignore_same_index) {
   CHECK_CUDA(x);
diff --git a/csrc/cuda/radius_cuda.h b/csrc/cuda/radius_cuda.h
index 7bf8908..d0a8282 100644
--- a/csrc/cuda/radius_cuda.h
+++ b/csrc/cuda/radius_cuda.h
@@ -3,7 +3,7 @@
 #include "../extensions.h"
 
 torch::Tensor radius_cuda(torch::Tensor x, torch::Tensor y,
-                          torch::optional<torch::Tensor> ptr_x,
-                          torch::optional<torch::Tensor> ptr_y, double r,
+                          std::optional<torch::Tensor> ptr_x,
+                          std::optional<torch::Tensor> ptr_y, double r,
                           int64_t max_num_neighbors,
                           bool ignore_same_index);
diff --git a/csrc/graclus.cpp b/csrc/graclus.cpp
index b986513..7333855 100644
--- a/csrc/graclus.cpp
+++ b/csrc/graclus.cpp
@@ -20,7 +20,7 @@ PyMODINIT_FUNC PyInit__graclus_cpu(void) { return NULL; }
 #endif
 
 CLUSTER_API torch::Tensor graclus(torch::Tensor rowptr, torch::Tensor col,
-                      torch::optional<torch::Tensor> optional_weight) {
+                      std::optional<torch::Tensor> optional_weight) {
   if (rowptr.device().is_cuda()) {
 #ifdef WITH_CUDA
     return graclus_cuda(rowptr, col, optional_weight);
diff --git a/csrc/grid.cpp b/csrc/grid.cpp
index 21b6688..f77512d 100644
--- a/csrc/grid.cpp
+++ b/csrc/grid.cpp
@@ -20,8 +20,8 @@ PyMODINIT_FUNC PyInit__grid_cpu(void) { return NULL; }
 #endif
 
 CLUSTER_API torch::Tensor grid(torch::Tensor pos, torch::Tensor size,
-                   torch::optional<torch::Tensor> optional_start,
-                   torch::optional<torch::Tensor> optional_end) {
+                   std::optional<torch::Tensor> optional_start,
+                   std::optional<torch::Tensor> optional_end) {
   if (pos.device().is_cuda()) {
 #ifdef WITH_CUDA
     return grid_cuda(pos, size, optional_start, optional_end);
diff --git a/csrc/knn.cpp b/csrc/knn.cpp
index 9d0589b..149bf5f 100644
--- a/csrc/knn.cpp
+++ b/csrc/knn.cpp
@@ -20,8 +20,8 @@ PyMODINIT_FUNC PyInit__knn_cpu(void) { return NULL; }
 #endif
 
 CLUSTER_API torch::Tensor knn(torch::Tensor x, torch::Tensor y,
-                  torch::optional<torch::Tensor> ptr_x,
-                  torch::optional<torch::Tensor> ptr_y, int64_t k, bool cosine,
+                  std::optional<torch::Tensor> ptr_x,
+                  std::optional<torch::Tensor> ptr_y, int64_t k, bool cosine,
                   int64_t num_workers) {
   if (x.device().is_cuda()) {
 #ifdef WITH_CUDA
diff --git a/csrc/radius.cpp b/csrc/radius.cpp
index 27a588a..1ab8884 100644
--- a/csrc/radius.cpp
+++ b/csrc/radius.cpp
@@ -20,8 +20,8 @@ PyMODINIT_FUNC PyInit__radius_cpu(void) { return NULL; }
 #endif
 
 CLUSTER_API torch::Tensor radius(torch::Tensor x, torch::Tensor y,
-                     torch::optional<torch::Tensor> ptr_x,
-                     torch::optional<torch::Tensor> ptr_y, double r,
+                     std::optional<torch::Tensor> ptr_x,
+                     std::optional<torch::Tensor> ptr_y, double r,
                      int64_t max_num_neighbors, int64_t num_workers,
                      bool ignore_same_index) {
   if (x.device().is_cuda()) {