add long type for jagged op. (pytorch#1214)

Summary: Pull Request resolved: pytorch#1214 as title. Reviewed By: jiaqizhai, brad-mengchi, mjanderson09 Differential Revision: D37978119 fbshipit-source-id: c2c004dfb2e1483f6fbf6a415c9dd58d95599cb0
williamwen42 · Jul 22, 2022 · 49061a2 · 49061a2
1 parent 0551586
commit 49061a2
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/fbgemm_gpu/src/jagged_tensor_ops.cu b/fbgemm_gpu/src/jagged_tensor_ops.cu
@@ -667,8 +667,12 @@ class DenseToJaggedGPUOp
     at::cuda::OptionalCUDAGuard device_guard;
     device_guard.set_index(dense.get_device());
 
-    AT_DISPATCH_FLOATING_TYPES_AND_HALF(
-        values.scalar_type(), "jagged_dense_add_forward", [&] {
+    AT_DISPATCH_FLOATING_TYPES_AND2(
+        at::ScalarType::Half,
+        at::ScalarType::Long,
+        values.scalar_type(),
+        "jagged_dense_add_forward",
+        [&] {
           jagged_dense_elementwise_jagged_output_<scalar_t>(
               values,
               offsets,