Use TORCH_BOX in rnnt

Author: pearu

src/libtorchaudio/overdrive.cpp

@@ -12,7 +12,7 @@ template <typename T, size_t N>
 using TensorAccessor = torch::headeronly::HeaderOnlyTensorAccessor<T, N>;
 
 // TODO: eliminate accessor<T, N>(t) in favor of t.accessor<T, N>
-// after Tensor::accessor is supported in stable ABI
+// after Tensor::accessor is supported in stable ABI.
 template <typename T, size_t N>
 inline TensorAccessor<T, N> accessor(Tensor t) {
   return TensorAccessor<T, N>(

src/libtorchaudio/rnnt/gpu/compute.cu

@@ -14,10 +14,10 @@ using torch::headeronly::ScalarType;
 
 // Entry point into RNNT Loss
 std::tuple<Tensor, Tensor> compute(
-    const Tensor& logits,
-    const Tensor& targets,
-    const Tensor& logit_lengths,
-    const Tensor& target_lengths,
+    Tensor logits,
+    Tensor targets,
+    Tensor logit_lengths,
+    Tensor target_lengths,
     int64_t blank,
     double clamp,
     bool fused_log_softmax = true) {
@@ -148,23 +148,13 @@ std::tuple<Tensor, Tensor> compute(
   return std::make_tuple(costs, gradients);
 }
 
-void boxed_rnnt_loss(StableIValue* stack, uint64_t num_args, uint64_t num_outputs) {
-  STD_TORCH_CHECK(num_args == 7, "num_args must be 7");
-  STD_TORCH_CHECK(num_outputs == 2, "num_outputs must be 2");
-  std::tuple<Tensor, Tensor> res = compute(
-      /*logits*/torch::stable::detail::to<Tensor>(stack[0]),
-      /*targets*/torch::stable::detail::to<Tensor>(stack[1]),
-      /*logit_lengths*/torch::stable::detail::to<Tensor>(stack[2]),
-      /*target_lengths*/torch::stable::detail::to<Tensor>(stack[3]),
-      /*blank*/float(torch::stable::detail::to<int64_t>(stack[4])),
-      /*clamp*/torch::stable::detail::to<double>(stack[5]),
-      /*fused_log_softmax*/torch::stable::detail::to<bool>(stack[6]));
-  stack[0] = torch::stable::detail::from(std::get<0>(res));
-  stack[1] = torch::stable::detail::from(std::get<1>(res));
+STABLE_TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
+  m.def(
+      "rnnt_loss_forward(Tensor logits, Tensor targets, Tensor logit_lengths, Tensor target_lengths, int blank, double clamp, bool fused_log_softmax) -> (Tensor, Tensor)");
 }
 
 STABLE_TORCH_LIBRARY_IMPL(torchaudio, CUDA, m) {
-  m.impl("rnnt_loss_forward", &boxed_rnnt_loss);
+  m.impl("rnnt_loss_forward", TORCH_BOX(&compute));
 }
 
 } // namespace gpu