Fix CUDA dispatch (pytorch#4075)

spcyppt · facebook-github-bot · commit 0c177e90f0fe · 2025-05-05T15:15:58.000-07:00
Summary: Pull Request resolved: pytorch#4075 X-link: facebookresearch/FBGEMM#1159 https://fb.workplace.com/groups/fbgemmusers/permalink/9486923244721705/ Reviewed By: q10 Differential Revision: D73451924 fbshipit-source-id: 2c8892235f44f4e5d796bd9c4350583919f6b34c
diff --git a/fbgemm_gpu/src/jagged_tensor_ops/keyed_jagged_index_select_dim1.cu b/fbgemm_gpu/src/jagged_tensor_ops/keyed_jagged_index_select_dim1.cu
@@ -551,4 +551,7 @@ TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
   DISPATCH_TO_AUTOGRAD_CUDA(
       "keyed_jagged_index_select_dim1",
       fbgemm_gpu::keyed_jagged_index_select_dim_1_gpu);
+  DISPATCH_TO_CUDA(
+      "keyed_jagged_index_select_dim1",
+      fbgemm_gpu::keyed_jagged_index_select_dim_1_gpu);
 }

Original file line number	Diff line number	Diff line change
`@@ -551,4 +551,7 @@ TORCH_LIBRARY_FRAGMENT(fbgemm, m) {`
`551`	`551`	`DISPATCH_TO_AUTOGRAD_CUDA(`
`552`	`552`	`"keyed_jagged_index_select_dim1",`
`553`	`553`	`fbgemm_gpu::keyed_jagged_index_select_dim_1_gpu);`
	`554`	`+ DISPATCH_TO_CUDA(`
	`555`	`+ "keyed_jagged_index_select_dim1",`
	`556`	`+ fbgemm_gpu::keyed_jagged_index_select_dim_1_gpu);`
`554`	`557`	`}`