Use TORCH_BOX in forced_align/cpu.

Author: pearu

src/libtorchaudio/forced_align/compute.cpp

@@ -2,5 +2,9 @@
 
 STABLE_TORCH_LIBRARY_FRAGMENT(torchaudio, m) {
   m.def(
-      "forced_align(Tensor log_probs, Tensor targets, Tensor input_lengths, Tensor target_lengths, int blank) -> (Tensor, Tensor)");
+      "forced_align(Tensor log_probs,"
+      "Tensor targets,"
+      "Tensor input_lengths,"
+      "Tensor target_lengths,"
+      "int blank) -> (Tensor, Tensor)");
 }

src/libtorchaudio/forced_align/cpu/compute.cpp

@@ -147,10 +147,10 @@ template <typename scalar_t>
 const auto forced_align_int_impl = forced_align_impl<scalar_t, ScalarType::Int>;
 
 std::tuple<Tensor, Tensor> compute(
-    const Tensor& logProbs,
-    const Tensor& targets,
-    const Tensor& inputLengths,
-    const Tensor& targetLengths,
+    Tensor logProbs,
+    Tensor targets,
+    Tensor inputLengths,
+    Tensor targetLengths,
     const int64_t blank) {
   STD_TORCH_CHECK(logProbs.is_cpu(), "log_probs must be a CPU tensor");
   STD_TORCH_CHECK(targets.is_cpu(), "targets must be a CPU tensor");
@@ -224,24 +224,8 @@ std::tuple<Tensor, Tensor> compute(
   return std::make_tuple(paths, logProbs);
 }
 
-void boxed_forced_align_cpu(
-    StableIValue* stack,
-    uint64_t num_args,
-    uint64_t num_outputs) {
-  STD_TORCH_CHECK(num_args == 5, "num_args must be 5");
-  STD_TORCH_CHECK(num_outputs == 2, "num_outputs must be 2");
-  std::tuple<Tensor, Tensor> res = compute(
-      /*logProbs*/ 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])));
-  stack[0] = torch::stable::detail::from(std::get<0>(res));
-  stack[1] = torch::stable::detail::from(std::get<1>(res));
-}
-
 STABLE_TORCH_LIBRARY_IMPL(torchaudio, CPU, m) {
-  m.impl("forced_align", &boxed_forced_align_cpu);
+  m.impl("forced_align", TORCH_BOX(&compute));
 }
 
 } // namespace cpu

src/libtorchaudio/forced_align/gpu/compute.cu

@@ -318,13 +318,19 @@ std::tuple<Tensor, Tensor> compute(
   std::cout << "forced_align: compute: 2" << std::endl;
   THO_DISPATCH_V2(logProbs.scalar_type(), "forced_align_impl", AT_WRAP([&] {
         if (targets.scalar_type() == ScalarType::Long) {
+          std::cout << "forced_align: compute: 2.1" << std::endl;
           forced_align_long_impl<scalar_t>(logProbs, targets, blank, paths);
+          std::cout << "forced_align: compute: 2.2" << std::endl;
         } else {
+          std::cout << "forced_align: compute: 2.3" << std::endl;
+          STD_TORCH_CHECK(targets.scalar_type() == ScalarType::Int, "unexpected dtype");
           forced_align_int_impl<scalar_t>(logProbs, targets, blank, paths);
+          std::cout << "forced_align: compute: 2.4" << std::endl;
         }
       }), AT_EXPAND(AT_FLOATING_TYPES), ScalarType::Half);
-
+  std::cout << "forced_align: compute: 3" << std::endl;
   Tensor pathsCuda = torchaudio::stable::cuda(paths, logProbs.get_device_index());
+  std::cout << "forced_align: compute: 4" << std::endl;
   return std::make_tuple(pathsCuda, logProbs);
 }