Update embedding_split_host_pt2_autograd_template.cpp

spcyppt · Jan 31, 2025 · 14c8cc1 · 14c8cc1
1 parent c0efcc0
commit 14c8cc1
Showing 1 changed file with 2 additions and 9 deletions.
diff --git a/fbgemm_gpu/codegen/training/pt2/embedding_split_host_pt2_autograd_template.cpp b/fbgemm_gpu/codegen/training/pt2/embedding_split_host_pt2_autograd_template.cpp
@@ -148,7 +148,7 @@ enum SSDTensor {
       {%- endif %}
       hash_size_cumsum,
       indices,
-      offsets_,
+      offsets,
       {%- if not nobag %}
       pooling_mode,
       indice_weights_value,
@@ -706,13 +706,6 @@ class {{ autograd_func }} :
         info_B_num_bits,
         /*total_B=*/offsets.sym_size(0) - 1
         );
-      Tensor offsets_;
-      if (weights_host.numel()){
-          offsets_ = reshape_offsets(offsets_, B_offsets, max_B, T);
-      }
-      else {
-          offsets_ = offsets;
-      }
     {%- endif %} // vbe
 
     {%- if is_gwd %}
@@ -735,7 +728,7 @@ class {{ autograd_func }} :
         {%- endif %}
         hash_size_cumsum,
         indices,
-        offsets_,
+        offsets,
         {%- if not nobag %}
         indice_weights_value,
         feature_requires_grad.value_or(Tensor()),