Add length check

Author: samanklesaria

src/libtorchaudio/rnnt/cpu/compute.cpp

@@ -86,10 +86,13 @@ std::tuple<Tensor, Tensor> compute(
   aoti_torch_get_size(targets.get(), 0, &targets_size);
   AOTI_TORCH_CHECK(targets_size == logits_size);
 
-  // TORCH_CHECK(
-  //     blank >= 0 && blank < logits.size(-1),
-  //     "blank must be within [0, logits.shape[-1])");
+  AOTI_TORCH_CHECK(
+      blank >= 0 && blank < logits.size(-1),
+      "blank must be within [0, logits.shape[-1])");
 
+  // "Max" is not ABI stable yet, but no tests check
+  // for this error behavior, so it's okay to merge in for now.
+  //
   // TORCH_CHECK(
   //     logits.size(1) == at::max(logit_lengths).item().toInt(),
   //     "input length mismatch");

src/libtorchaudio/rnnt/gpu/compute.cu

@@ -84,9 +84,9 @@ std::tuple<Tensor, Tensor> compute(
     TORCH_ERROR_CODE_CHECK(aoti_torch_get_size(targets.get(), 0, &targets_size));
     AOTI_TORCH_CHECK(targets_size == logits_size);
 
-    // TORCH_CHECK(
-    //     blank >= 0 && blank < logits.size(-1),
-    //     "blank must be within [0, logits.shape[-1])");
+    AOTI_TORCH_CHECK(
+        blank >= 0 && blank < logits.size(-1),
+        "blank must be within [0, logits.shape[-1])");
 
     // TORCH_CHECK(
     //     logits.size(1) == at::max(logit_lengths).item().toInt(),