[MPS] Enable bucketization for bfloat16 (#136983)

malfet · AnantGulati · commit 9845028da443 · 2024-10-02T14:27:03.000+03:00
By simply adding explicit instantiation Tested in #136987 Pull Request resolved: #136983 Approved by: https://github.com/Skylion007 ghstack dependencies: #136981, #136982
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);
