diff --git a/hopper/flash_api.cpp b/hopper/flash_api.cpp
index 5921c374f6b..5b3d124627a 100644
--- a/hopper/flash_api.cpp
+++ b/hopper/flash_api.cpp
@@ -3,8 +3,8 @@
  ******************************************************************************/
 
 #include <Python.h>
-#include <torch/nn/functional.h>
-#include <ATen/cuda/CUDAContext.h>
+#include <torch/nn/functional/padding.h>
+#include <ATen/cuda/CUDAContextLight.h>
 #include <c10/cuda/CUDAGuard.h>
 
 #include <cutlass/numeric_types.h>