SymInt-ify scheme batch index select (pytorch#2488)

Summary:
Pull Request resolved: pytorch#2488

SymInt-ify scheme of batch_index_select as it can be used via torchrec VB path with SymInt

Reviewed By: ezyang

Differential Revision: D55923134

fbshipit-source-id: b7b69d307c866cf39fad3e6e44a50e24295a0109

fbgemm_gpu/codegen/training/index_select/batch_index_select_dim0_cpu_host.cpp

@@ -217,9 +217,9 @@ TORCH_LIBRARY_FRAGMENT(fb, m) {
       "batch_index_select_dim0("
       "    Tensor inputs,"
       "    Tensor indices,"
-      "    int[] input_num_indices,"
-      "    int[] input_rows,"
-      "    int[] input_columns,"
+      "    SymInt[] input_num_indices,"
+      "    SymInt[] input_rows,"
+      "    SymInt[] input_columns,"
       "    bool permute_output_dim_0_1=False) -> Tensor");
   DISPATCH_TO_CPU("batch_index_select_dim0", batch_index_select_dim0_cpu);
 }
@@ -232,9 +232,9 @@ TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
       "batch_index_select_dim0("
       "    Tensor inputs,"
       "    Tensor indices,"
-      "    int[] input_num_indices,"
-      "    int[] input_rows,"
-      "    int[] input_columns,"
+      "    SymInt[] input_num_indices,"
+      "    SymInt[] input_rows,"
+      "    SymInt[] input_columns,"
       "    bool permute_output_dim_0_1=False) -> Tensor");
   DISPATCH_TO_CPU("batch_index_select_dim0", batch_index_select_dim0_cpu);
 }