Fix prev_iter (#3784)

Summary: X-link: facebookresearch/FBGEMM#868 Pull Request resolved: #3784 as titled. Reviewed By: q10, sryap Differential Revision: D70794993 fbshipit-source-id: 07004d78a6112c972d77a3fdf0c74519dcd4019b
pytorch · Mar 7, 2025 · 1b31500 · 1b31500
1 parent eb9108c
commit 1b31500
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/fbgemm_gpu/codegen/training/python/split_embedding_codegen_lookup_invoker.template b/fbgemm_gpu/codegen/training/python/split_embedding_codegen_lookup_invoker.template
@@ -213,7 +213,7 @@ def invoke(
     {%- if "prev_iter" not in args_pt2.unified_pt2.split_unpacked_arg_names %}
     dict_aux_tensor["prev_iter_dev"] = prev_iter_dev
     {%- else %}
-    dict_aux_tensor["prev_iter_dev"] = prev_iter.prev_iter_dev
+    dict_aux_tensor["prev_iter_dev"] = prev_iter.dev
     {%- endif %}
 
     # optimizer_args