diff --git a/functorch/csrc/BatchRulesViews.cpp b/functorch/csrc/BatchRulesViews.cpp index de4f3830f..abd8e2602 100644 --- a/functorch/csrc/BatchRulesViews.cpp +++ b/functorch/csrc/BatchRulesViews.cpp @@ -479,6 +479,7 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) { VMAP_SUPPORT("expand", expand_batch_rule); VMAP_SUPPORT("unfold", unfold_batch_rule); VMAP_SUPPORT("movedim.intlist", movedim_batch_rule); + m.impl("movedim.int", static_cast(native::movedim)); // composite wrt autograd } }}