readd movedim

pytorch · Oct 23, 2021 · c9a3419 · c9a3419
1 parent ab4a32f
commit c9a3419
Showing 1 changed file with 1 addition and 0 deletions.
diff --git 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<Tensor(*)(const Tensor&,int64_t,int64_t)>(native::movedim)); // composite wrt autograd
 } 
 
 }}