diff --git a/CMakeLists.txt b/CMakeLists.txt
index 78ee10e..aed35dd 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -32,18 +32,18 @@ if(NOT CMAKE_BUILD_TYPE)
   set(CMAKE_BUILD_TYPE Release)
 endif()
 
-set(CMAKE_CXX_STANDARD 14 CACHE STRING "The C++ version to be used.")
-set(CMAKE_CXX_EXTENSIONS OFF)
-
 message(STATUS "C++ Standard version: ${CMAKE_CXX_STANDARD}")
 
 list(APPEND CMAKE_MODULE_PATH ${CMAKE_CURRENT_SOURCE_DIR}/cmake/Modules)
 list(APPEND CMAKE_MODULE_PATH ${CMAKE_CURRENT_SOURCE_DIR}/cmake)
 
+option(CMAKE_CXX_STANDARD 14 "The C++ version to be used." 14)
 option(BUILD_SHARED_LIBS "Whether to build shared libraries" ON)
 option(kaldifeat_BUILD_TESTS "Whether to build tests or not" ON)
 option(kaldifeat_BUILD_PYMODULE "Whether to build python module or not" ON)
 
+set(CMAKE_CXX_EXTENSIONS OFF)
+
 message(STATUS "BUILD_SHARED_LIBS ${BUILD_SHARED_LIBS}")
 
 if(BUILD_SHARED_LIBS AND MSVC)
diff --git a/cmake/cmake_extension.py b/cmake/cmake_extension.py
index 618b289..819c0eb 100644
--- a/cmake/cmake_extension.py
+++ b/cmake/cmake_extension.py
@@ -8,9 +8,15 @@ import sys
 from pathlib import Path
 
 import setuptools
+import torch
 from setuptools.command.build_ext import build_ext
 
 
+def get_pytorch_version():
+    # if it is 1.7.1+cuda101, then strip +cuda101
+    return torch.__version__.split("+")[0]
+
+
 def is_for_pypi():
     ans = os.environ.get("KALDIFEAT_IS_FOR_PYPI", None)
     return ans is not None
@@ -39,7 +45,6 @@ try:
                 # -linux_x86_64.whl
                 self.root_is_pure = False
 
-
 except ImportError:
     bdist_wheel = None
 
@@ -70,6 +75,12 @@ class BuildExtension(build_ext):
         extra_cmake_args = " -Dkaldifeat_BUILD_TESTS=OFF "
         extra_cmake_args += f" -DCMAKE_INSTALL_PREFIX={Path(self.build_lib).resolve()}/kaldifeat "  # noqa
 
+        major, minor = get_pytorch_version().split(".")[:2]
+        major = int(major)
+        minor = int(minor)
+        if major > 2 or (major == 2 and minor >= 1):
+            extra_cmake_args += f" -DCMAKE_CXX_STANDARD=17 "
+
         if "PYTHON_EXECUTABLE" not in cmake_args:
             print(f"Setting PYTHON_EXECUTABLE to {sys.executable}")
             cmake_args += f" -DPYTHON_EXECUTABLE={sys.executable}"
@@ -103,9 +114,7 @@ class BuildExtension(build_ext):
         else:
             if make_args == "" and system_make_args == "":
                 print("For fast compilation, run:")
-                print(
-                    'export KALDIFEAT_MAKE_ARGS="-j"; python setup.py install'
-                )
+                print('export KALDIFEAT_MAKE_ARGS="-j"; python setup.py install')
                 make_args = " -j4 "
                 print("Setting make_args to '-j4'")
 
diff --git a/setup.py b/setup.py
index 2271dd0..123ca5f 100644
--- a/setup.py
+++ b/setup.py
@@ -9,6 +9,7 @@ import setuptools
 from cmake.cmake_extension import BuildExtension, bdist_wheel, cmake_extension
 import get_version
 
+
 get_package_version = get_version.get_package_version