use `torch.bucketize` in`to_sparse_csr` implementation (+ additional tests) (#61340)
Summary:
Fixes https://github.com/pytorch/pytorch/issues/57381
Pull Request resolved: https://github.com/pytorch/pytorch/pull/61340
Reviewed By: bhosmer
Differential Revision: D29601393
Pulled By: cpuhrsch
fbshipit-source-id: 4ca1f013d96e8716f0e658e0cd685d9aa0d98a5c
