diff --git a/fbgemm_gpu/include/fbgemm_gpu/permute_pooled_embedding_ops.h b/fbgemm_gpu/include/fbgemm_gpu/permute_pooled_embedding_ops.h
index 1e08490ed..9c1759772 100644
--- a/fbgemm_gpu/include/fbgemm_gpu/permute_pooled_embedding_ops.h
+++ b/fbgemm_gpu/include/fbgemm_gpu/permute_pooled_embedding_ops.h
@@ -9,7 +9,8 @@
 #pragma once
 
 #include <ATen/ATen.h>
-#include <torch/script.h>
+#include <torch/csrc/api/include/torch/types.h>
+#include <torch/csrc/autograd/custom_function.h>
 #include "fbgemm_gpu/ops_utils.h"
 #include "fbgemm_gpu/sparse_ops_utils.h"