Skip to content

Commit

Permalink
RW Dist change to support uneven sharding (pytorch#2142)
Browse files Browse the repository at this point in the history
Summary:
X-link: pytorch/torchrec#1525

Pull Request resolved: pytorch#2142

As titled

Reviewed By: dstaay-fb

Differential Revision: D51053207

fbshipit-source-id: 8aee7d967ceec9ea0739f5d2c56d6e541a0d3648
  • Loading branch information
gnahzg authored and facebook-github-bot committed Nov 18, 2023
1 parent ad2aca2 commit 92388c1
Showing 1 changed file with 2 additions and 1 deletion.
3 changes: 2 additions & 1 deletion fbgemm_gpu/src/sparse_ops/sparse_ops_cpu.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -306,7 +306,8 @@ void _block_bucketize_sparse_features_cpu(
const index_t* const block_sizes_data = block_sizes.data_ptr<index_t>();
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<index_t>;
using uoffset_t = std::make_unsigned_t<offset_t>;
std::vector<int64_t> lower_bounds(indices.numel(), 0);
Expand Down

0 comments on commit 92388c1

Please sign in to comment.