embedding decomp

pytorch · Jun 16, 2022 · 3b68942 · 3b68942
1 parent 825f439
commit 3b68942
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/functorch/_src/eager_transforms.py b/functorch/_src/eager_transforms.py
@@ -1339,7 +1339,6 @@ def _register_python_decomposition_vmap(decomp):
 _register_jit_decomposition(torch.ops.aten._softmax_backward_data.default)
 _register_jit_decomposition(torch.ops.aten.log_sigmoid_forward.default)
 _register_jit_decomposition(torch.ops.aten.native_layer_norm_backward.default)
-_register_jit_decomposition(torch.ops.aten.native_batch_norm_backward.default)
-_register_jit_decomposition(torch.ops.aten.cudnn_batch_norm_backward.default)
+_register_jit_decomposition(torch.ops.aten.embedding_dense_backward.default)
 _register_python_decomposition_vmap(torch.ops.aten.mse_loss_backward.default)
 _register_python_decomposition_vmap(torch.ops.aten.addr.default)
diff --git a/functorch/csrc/DynamicLayer.cpp b/functorch/csrc/DynamicLayer.cpp
@@ -504,6 +504,7 @@ TORCH_LIBRARY_IMPL(aten, FT_DYNAMIC_LAYER_FRONT_MODE_KEY, m) {
   JVP_DECOMP(native_layer_norm_backward);
   JVP_DECOMP(native_batch_norm_backward);
   JVP_DECOMP(cudnn_batch_norm_backward);
+  JVP_DECOMP(embedding_dense_backward);
 }