diff --git a/paddle/pten/backends/gpu/gpu_launch_config.h b/paddle/pten/backends/gpu/gpu_launch_config.h
index c166b3d2f80ed3075594455d6d95df4866503f84..edb35ef4ecc6e94dafa2e63fb19a4daa76bf8fac 100644
--- a/paddle/pten/backends/gpu/gpu_launch_config.h
+++ b/paddle/pten/backends/gpu/gpu_launch_config.h
@@ -64,7 +64,7 @@ static inline int RoundToPowerOfTwo(int n) {
 #ifdef WITH_NV_JETSON
 // The number of threads cannot be assigned 1024 in some cases when the device
 // is nano or tx2 .
-inline void ChangeThreadNum(const pten::CUDAContext& context,
+inline void ChangeThreadNum(const pten::GPUContext& context,
                             int* num_thread,
                             int alternative_num_thread = 512) {
   if (context.GetComputeCapability() == 53 ||