diff --git a/python/aitemplate/backend/cuda/conv2d/common.py b/python/aitemplate/backend/cuda/conv2d/common.py
index c6bfbdb8b..76822e244 100644
--- a/python/aitemplate/backend/cuda/conv2d/common.py
+++ b/python/aitemplate/backend/cuda/conv2d/common.py
@@ -364,11 +364,14 @@
 
 PROFILER_MAIN_TEMPLATE = jinja2.Template(
     """
+#include <cstdint>
 #include <iostream>
 #include <string>
 
 #include "cutlass/cutlass.h"
 
+using std::int8_t;
+
 {{benchmark_decls}}
 
 int main(int argc, char** argv) {