Update embedding_split_host_pt2_cpu_wrapper_template.cpp

spcyppt · Jan 31, 2025 · 54cdaf7 · 54cdaf7
1 parent ca3ebf4
commit 54cdaf7
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/fbgemm_gpu/codegen/training/pt2/embedding_split_host_pt2_cpu_wrapper_template.cpp b/fbgemm_gpu/codegen/training/pt2/embedding_split_host_pt2_cpu_wrapper_template.cpp
@@ -27,6 +27,10 @@
 #include "fbgemm_gpu/utils/dispatch_macros.h"
 #include "fbgemm_gpu/embedding_common.h"
 
+{%- if has_vbe_support %}
+#include "fbgemm_gpu/utils/pt2_autograd_utils.h"
+{%- endif %}
+
 using Tensor = at::Tensor;
 using namespace fbgemm_gpu;