diff --git a/kaldifeat/csrc/test_kaldifeat.cc b/kaldifeat/csrc/test_kaldifeat.cc
index 0c56056..f35d38e 100644
--- a/kaldifeat/csrc/test_kaldifeat.cc
+++ b/kaldifeat/csrc/test_kaldifeat.cc
@@ -2,6 +2,7 @@
 //
 // Copyright (c)  2021  Xiaomi Corporation (authors: Fangjun Kuang)
 
+#include "torch/all.h"
 #include "torch/script.h"
 
 static void TestPreemph() {
diff --git a/kaldifeat/python/csrc/kaldifeat.h b/kaldifeat/python/csrc/kaldifeat.h
index 11865e5..0c39d3d 100644
--- a/kaldifeat/python/csrc/kaldifeat.h
+++ b/kaldifeat/python/csrc/kaldifeat.h
@@ -6,6 +6,7 @@
 #define KALDIFEAT_PYTHON_CSRC_KALDIFEAT_H_
 
 #include "pybind11/pybind11.h"
+#include "torch/torch.h"
 namespace py = pybind11;
 
 #endif  // KALDIFEAT_PYTHON_CSRC_KALDIFEAT_H_