From 54cdaf76b9166ceb6fb0cb6e37f64096939f9be3 Mon Sep 17 00:00:00 2001
From: Supadchaya <138070207+spcyppt@users.noreply.github.com>
Date: Fri, 31 Jan 2025 00:06:25 -0800
Subject: [PATCH] Update embedding_split_host_pt2_cpu_wrapper_template.cpp

---
 .../pt2/embedding_split_host_pt2_cpu_wrapper_template.cpp     | 4 ++++
 1 file changed, 4 insertions(+)

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
index 1c25ae326..57275a455 100644
--- 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;