From 92388c170c0b9b8318fbecfcad50b75a666b923f Mon Sep 17 00:00:00 2001 From: Qiang Zhang Date: Sat, 18 Nov 2023 04:59:48 -0800 Subject: [PATCH] RW Dist change to support uneven sharding (#2142) Summary: X-link: https://github.com/pytorch/torchrec/pull/1525 Pull Request resolved: https://github.com/pytorch/FBGEMM/pull/2142 As titled Reviewed By: dstaay-fb Differential Revision: D51053207 fbshipit-source-id: 8aee7d967ceec9ea0739f5d2c56d6e541a0d3648 --- fbgemm_gpu/src/sparse_ops/sparse_ops_cpu.cpp | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/fbgemm_gpu/src/sparse_ops/sparse_ops_cpu.cpp b/fbgemm_gpu/src/sparse_ops/sparse_ops_cpu.cpp index 8291927f45..7acefdf18a 100644 --- a/fbgemm_gpu/src/sparse_ops/sparse_ops_cpu.cpp +++ b/fbgemm_gpu/src/sparse_ops/sparse_ops_cpu.cpp @@ -306,7 +306,8 @@ void _block_bucketize_sparse_features_cpu( const index_t* const block_sizes_data = block_sizes.data_ptr(); offset_t* batch_sizes_data = nullptr; const auto variable_batch_size = batch_size_per_feature.has_value(); - const auto variable_bucket_sizes = block_bucketize_pos.has_value(); + const auto variable_bucket_sizes = block_bucketize_pos.has_value() && + block_bucketize_pos.value().size() != 0; using uindex_t = std::make_unsigned_t; using uoffset_t = std::make_unsigned_t; std::vector lower_bounds(indices.numel(), 0);