[MPS] Enable bucketization for bfloat16

malfet · malfet · commit ed659b1ad4d2 · 2024-09-29T22:29:21.000-07:00
[ghstack-poisoned]
diff --git a/aten/src/ATen/native/mps/operations/Bucketization.mm b/aten/src/ATen/native/mps/operations/Bucketization.mm
@@ -183,6 +183,10 @@ kernel void searchsorted(
 REGISTER_SEARCHSORTED_OP(float, long);
 REGISTER_SEARCHSORTED_OP(half, int);
 REGISTER_SEARCHSORTED_OP(half, long);
+#if __METAL_VERSION__ >= 310
+REGISTER_SEARCHSORTED_OP(bfloat, int);
+REGISTER_SEARCHSORTED_OP(bfloat, long);
+#endif
 REGISTER_SEARCHSORTED_OP(char, int);
 REGISTER_SEARCHSORTED_OP(char, long);
 REGISTER_SEARCHSORTED_OP(uchar, int);
