[gorkadabo1]

What does this PR do?

Fixes #4410

This PR computes the final-batch remainder using the sized sampler, following batch-sampler wrappers and respecting drop_last. For unsized samplers, it keeps the existing dataset-length fallback. It also sets the shard sampler's epoch before reading its length, so metric gathering stays correct when the sampler length changes between epochs.

Related PRs

I also reviewed #4412 (closed) and #4489 (open), which aim to fix the same issue. Both pass the original SubsetRandomSampler reproduction, but I found additional cases where they still return incorrect results:

• Distributed skip_first_batches: with zero batches skipped, both return 13 samples instead of 10. With one batch skipped, they return 5 instead of 2.
• Nested/top-level batch samplers: both return 13 samples instead of 10 in the tested two-process case.
• Changing sampler lengths: Fix gather_for_metrics dropping or duplicating samples from sized samplers #4489 returns 5 instead of 6 samples in the test, while Fix gather_for_metrics remainder when a sized sampler selects a subset #4412 returns the correct number.

This PR also handles these cases with a smaller change to the production code.

Validation

• Both new regression tests fail on upstream main and pass with the fix.
• Original CPU reproduction: 10/10 samples with one and two processes.
• Relevant tests: 88 passed, 16 skipped.
• I also tested the cases above on CPU, along with unsized samplers and explicit drop_last.
• make style, make quality and git diff --check pass.

I did not run the full test suite or GPU tests.

Before submitting

• This PR fixes a typo or improves the docs (not applicable).
• Did you read the contributor guideline, Pull Request section?
• Was this discussed/approved via a GitHub issue or the forum? Discussed in gather_for_metrics drops or duplicates examples with SubsetRandomSampler #4410.
• Did you make sure to update the documentation with your changes? (Not applicable: no public API change.)
• Did you write any new necessary tests?

Who can review?

Core library: @BenjaminBossan @SunMarc

[JonnyFaeth]

Additional coverage at 18954ea8eb4786b4db10f8aba7ae1bdf00b4a8f6: a WeightedRandomSampler draws 14 samples with replacement from a 5-item dataset, through both ordinary batching and an explicit BatchSampler.

Unlike a subset, oversampling produces intentional repeats that metrics must retain. The tests compare the complete gathered sequence with a separately seeded sampler and assert a remainder of 2. Both pass on this PR; restoring the parent data_loader.py makes both fail because the final sample is lost (dataset length gives remainder 1 instead of 2).

Validation: 4 passed, 40 deselected (2 new + the PR's subset and changing-length regressions); Ruff check/format passed. This is real single-process CPU Accelerator.prepare() / gather_for_metrics() on Windows/Python 3.12.14. It does not add distributed, XLA, or GPU validation, and I did not repeat the multi-process cases already described in your PR.

Prepared and tested with OpenAI Codex assistance.

Optional regression patch

diff --git a/tests/test_weighted_metrics.py b/tests/test_weighted_metrics.py
new file mode 100644
--- /dev/null
+++ b/tests/test_weighted_metrics.py
@@ -0,0 +1,42 @@
+# Copyright 2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+import torch
+from torch.utils.data import BatchSampler, DataLoader, WeightedRandomSampler
+
+from accelerate import Accelerator
+
+
+@pytest.mark.parametrize("explicit_batch_sampler", [False, True])
+def test_metrics_keep_all_weighted_replacement_samples(explicit_batch_sampler):
+    weights = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0])
+
+    def make_sampler():
+        return WeightedRandomSampler(weights, 14, replacement=True, generator=torch.Generator().manual_seed(123))
+
+    expected = list(make_sampler())
+    sampler = make_sampler()
+    loader = (
+        DataLoader(range(5), batch_sampler=BatchSampler(sampler, 4, drop_last=False))
+        if explicit_batch_sampler
+        else DataLoader(range(5), batch_size=4, sampler=sampler)
+    )
+    accelerator = Accelerator(cpu=True)
+    prepared = accelerator.prepare(loader)
+    actual = []
+    for batch in prepared:
+        actual.extend(accelerator.gather_for_metrics(batch).tolist())
+    assert actual == expected
+    assert prepared.remainder == 2