diff --git a/kaldifeat/csrc/feature-fbank.cc b/kaldifeat/csrc/feature-fbank.cc
index e740201..9945541 100644
--- a/kaldifeat/csrc/feature-fbank.cc
+++ b/kaldifeat/csrc/feature-fbank.cc
@@ -8,8 +8,6 @@
 
 #include <cmath>
 
-#include "torch/torch.h"
-
 namespace kaldifeat {
 
 std::ostream &operator<<(std::ostream &os, const FbankOptions &opts) {
diff --git a/kaldifeat/csrc/feature-fbank.h b/kaldifeat/csrc/feature-fbank.h
index f964ba7..05a467f 100644
--- a/kaldifeat/csrc/feature-fbank.h
+++ b/kaldifeat/csrc/feature-fbank.h
@@ -13,7 +13,6 @@
 #include "kaldifeat/csrc/feature-common.h"
 #include "kaldifeat/csrc/feature-window.h"
 #include "kaldifeat/csrc/mel-computations.h"
-#include "torch/torch.h"
 
 namespace kaldifeat {
 
diff --git a/kaldifeat/csrc/feature-functions.h b/kaldifeat/csrc/feature-functions.h
index 7a9bffd..bdc3c2c 100644
--- a/kaldifeat/csrc/feature-functions.h
+++ b/kaldifeat/csrc/feature-functions.h
@@ -7,7 +7,7 @@
 #ifndef KALDIFEAT_CSRC_FEATURE_FUNCTIONS_H_
 #define KALDIFEAT_CSRC_FEATURE_FUNCTIONS_H_
 
-#include "torch/torch.h"
+#include "torch/script.h"
 
 namespace kaldifeat {
 
diff --git a/kaldifeat/csrc/feature-mfcc.h b/kaldifeat/csrc/feature-mfcc.h
index 88ba815..2c00b7f 100644
--- a/kaldifeat/csrc/feature-mfcc.h
+++ b/kaldifeat/csrc/feature-mfcc.h
@@ -13,7 +13,7 @@
 #include "kaldifeat/csrc/feature-common.h"
 #include "kaldifeat/csrc/feature-window.h"
 #include "kaldifeat/csrc/mel-computations.h"
-#include "torch/torch.h"
+#include "torch/script.h"
 
 namespace kaldifeat {
 
diff --git a/kaldifeat/csrc/feature-plp.cc b/kaldifeat/csrc/feature-plp.cc
index 6e0a077..5b3871c 100644
--- a/kaldifeat/csrc/feature-plp.cc
+++ b/kaldifeat/csrc/feature-plp.cc
@@ -7,7 +7,6 @@
 #include "kaldifeat/csrc/feature-plp.h"
 
 #include "kaldifeat/csrc/feature-functions.h"
-#include "torch/torch.h"
 
 namespace kaldifeat {
 
diff --git a/kaldifeat/csrc/feature-plp.h b/kaldifeat/csrc/feature-plp.h
index 7c29dfc..dd5247c 100644
--- a/kaldifeat/csrc/feature-plp.h
+++ b/kaldifeat/csrc/feature-plp.h
@@ -13,7 +13,7 @@
 #include "kaldifeat/csrc/feature-common.h"
 #include "kaldifeat/csrc/feature-window.h"
 #include "kaldifeat/csrc/mel-computations.h"
-#include "torch/torch.h"
+#include "torch/script.h"
 
 namespace kaldifeat {
 
diff --git a/kaldifeat/csrc/feature-spectrogram.h b/kaldifeat/csrc/feature-spectrogram.h
index f51d450..c4febe6 100644
--- a/kaldifeat/csrc/feature-spectrogram.h
+++ b/kaldifeat/csrc/feature-spectrogram.h
@@ -11,7 +11,7 @@
 
 #include "kaldifeat/csrc/feature-common.h"
 #include "kaldifeat/csrc/feature-window.h"
-#include "torch/torch.h"
+#include "torch/script.h"
 
 namespace kaldifeat {
 
diff --git a/kaldifeat/csrc/feature-window.cc b/kaldifeat/csrc/feature-window.cc
index 4d24c8e..4ec5ac2 100644
--- a/kaldifeat/csrc/feature-window.cc
+++ b/kaldifeat/csrc/feature-window.cc
@@ -9,8 +9,6 @@
 #include <cmath>
 #include <vector>
 
-#include "torch/torch.h"
-
 #ifndef M_2PI
 #define M_2PI 6.283185307179586476925286766559005
 #endif
diff --git a/kaldifeat/csrc/feature-window.h b/kaldifeat/csrc/feature-window.h
index 03248f0..255c274 100644
--- a/kaldifeat/csrc/feature-window.h
+++ b/kaldifeat/csrc/feature-window.h
@@ -7,7 +7,8 @@
 #include <string>
 
 #include "kaldifeat/csrc/log.h"
-#include "torch/torch.h"
+#include "torch/all.h"
+#include "torch/script.h"
 
 #ifndef KALDIFEAT_CSRC_FEATURE_WINDOW_H_
 #define KALDIFEAT_CSRC_FEATURE_WINDOW_H_
diff --git a/kaldifeat/csrc/matrix-functions.h b/kaldifeat/csrc/matrix-functions.h
index 7bbcdaf..c9e493b 100644
--- a/kaldifeat/csrc/matrix-functions.h
+++ b/kaldifeat/csrc/matrix-functions.h
@@ -7,7 +7,7 @@
 #ifndef KALDIFEAT_CSRC_MATRIX_FUNCTIONS_H_
 #define KALDIFEAT_CSRC_MATRIX_FUNCTIONS_H_
 
-#include "torch/torch.h"
+#include "torch/script.h"
 
 namespace kaldifeat {
 
diff --git a/kaldifeat/csrc/pitch-functions.h b/kaldifeat/csrc/pitch-functions.h
index b703feb..32b50dd 100644
--- a/kaldifeat/csrc/pitch-functions.h
+++ b/kaldifeat/csrc/pitch-functions.h
@@ -20,7 +20,7 @@
 
 #include <string>
 
-#include "torch/torch.h"
+#include "torch/script.h"
 
 namespace kaldifeat {
 
diff --git a/kaldifeat/csrc/test_kaldifeat.cc b/kaldifeat/csrc/test_kaldifeat.cc
index bd2d43e..0c56056 100644
--- a/kaldifeat/csrc/test_kaldifeat.cc
+++ b/kaldifeat/csrc/test_kaldifeat.cc
@@ -2,7 +2,7 @@
 //
 // Copyright (c)  2021  Xiaomi Corporation (authors: Fangjun Kuang)
 
-#include "torch/torch.h"
+#include "torch/script.h"
 
 static void TestPreemph() {
   torch::Tensor a = torch::arange(0, 12).reshape({3, 4}).to(torch::kFloat);