Fix to batching rule (#685)

The to(self, other) batching rule doesn't actually handle the case where
self is not Batched and other is Batched.

Author: zou3519

functorch/csrc/BatchRulesUnaryOps.cpp

@@ -71,6 +71,14 @@ view_as_complex_batch_rule(const Tensor& self, optional<int64_t> self_bdim) {
   return std::make_tuple(result, 0);
 }
 
+std::tuple<Tensor,optional<int64_t>>
+to_other_batch_rule(const Tensor& self, optional<int64_t> self_bdim,
+                    const Tensor& other, optional<int64_t> other_bdim,
+                    bool non_blocking,
+                    bool copy, c10::optional<at::MemoryFormat> memory_format) {
+  return std::make_tuple(self.to(other, non_blocking, copy, memory_format), self_bdim);
+}
+
 TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
 
 #define UNARY_POINTWISE_ALL2(op, overload) \
@@ -89,6 +97,7 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   VMAP_SUPPORT2(to, device, BASIC_UNARY_BATCH_RULE(ATEN_FN2(to, device)));
   VMAP_SUPPORT2(to, dtype, BASIC_UNARY_BATCH_RULE(ATEN_FN2(to, dtype)));
   VMAP_SUPPORT2(to, dtype_layout, BASIC_UNARY_BATCH_RULE(ATEN_FN2(to, dtype_layout)));
+  VMAP_SUPPORT2(to, other, to_other_batch_rule);
 
   UNARY_POINTWISE(_to_copy);
   UNARY_POINTWISE(alias);

functorch/csrc/BatchingRegistrations.cpp

@@ -672,15 +672,6 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   m.impl("squeeze_.dim", squeeze_dim__batching_rule);
   m.impl("unsqueeze_", unsqueeze__batching_rule);
 
-  // still legacy b/c this op is weird
-#define TO_BATCHING_RULE(name, ...) \
-  { \
-    using to_type = Tensor(Tensor::*)(__VA_ARGS__) const; \
-    m.impl(name, unwrap_and_call_method< \
-        to_type, &Tensor::to, __VA_ARGS__>);\
-  }
-  TO_BATCHING_RULE("to.other", const Tensor&, bool, bool, optional<MemoryFormat>)
-
   // still legacy because these are ridiculously complicated
   m.impl("as_strided", as_strided_batching_rule);
   m.impl("new_empty_strided", new_empty_strided_batching_rule);