huggingface/pytorch-image-models · warning

Rank {self.rank}: Assigned {scheduled_samples_count} samples

Error message

Rank {self.rank}: Assigned {scheduled_samples_count} samples to batches, but expected {effective_samples_this_rank} effective samples this epoch. Indices remaining: {effective_samples_this_rank - scheduled_samples_count}.

What it means

Final sanity check in NaFlexDataset._prepare_epoch_batches: after assigning indices to schedule batches, scheduled_samples_count should equal effective_samples_this_rank. A warning means some indices were never assigned to any batch (leftovers), so the epoch silently sees fewer samples than requested.

Source

Thrown at timm/data/naflex_dataset.py:480

            # Ensure we don't try to grab more indices than available for the rank
            actual_bs = min(bs, effective_samples_this_rank - idx_pos)
            if actual_bs <= 0:
                 if scheduled_samples_count < effective_samples_this_rank:
                     # This indicates mismatch between schedule total and actual samples
                     warnings.warn(f"Rank {self.rank}: Ran out of samples ({idx_pos}/{effective_samples_this_rank}) before processing entire schedule. Check schedule generation.")
                 break # Stop if no more indices or batch size is zero

            batch_indices = indices_this_rank[idx_pos : idx_pos + actual_bs]
            patch_idx = 0
            if self.variable_patch_size:
                patch_idx = torch.multinomial(patch_size_probs, 1, generator=g).item()
            epoch_batches.append((seq_len, patch_idx, batch_indices))
            idx_pos += actual_bs
            scheduled_samples_count += actual_bs

        # Final check
        if scheduled_samples_count != effective_samples_this_rank:
             warnings.warn(
                f"Rank {self.rank}: Assigned {scheduled_samples_count} samples to batches, "
                f"but expected {effective_samples_this_rank} effective samples this epoch. "
                f"Indices remaining: {effective_samples_this_rank - scheduled_samples_count}."
             )
        return epoch_batches

    def set_epoch(self, epoch: int) -> None:
        """Set the multiprocessing-safe epoch read by DataLoader workers.

        Args:
            epoch: New epoch number.
        """
        self.shared_epoch.value = epoch

    def __len__(self) -> int:
        """Return the number of batches for this rank.

        Returns:

View on GitHub (pinned to 9a5261e31b)

Solutions

1. Tune min_batch_size/batch_divisor so total scheduled == effective samples, or pad dataset accordingly
2. Verify world_size/rank consistency
3. If dropping a few samples per epoch is acceptable, suppress with warnings.filterwarnings for this message

Example fix

# before
NaFlexDataset(ds, ..., batch_divisor=32)  # 1000 samples/rank leaves remainder
# after
NaFlexDataset(ds, ..., batch_divisor=32, min_batch_size=1)  # remainder scheduled

assert sum(bs for _, bs in schedule) == effective_samples_per_rank, 'indices will go unused'

When it happens

Trigger: Batch_divisor/min_batch_size constraints leave a remainder that no schedule entry can hold (cascade from 267); clamping via min(bs, remaining) truncated batches; drop_last-style behavior on a non-divisible per-rank count.

Common situations: Same family as 267–269: odd dataset sizes under distributed sharding and packing constraints. Usually benign (a few dropped samples) but should be eliminated if exact epoch size matters.

Related errors

• Rank {self.rank}: Ran out of samples ({idx_pos}/{effective_s
• Rank {self.rank}: Canonical schedule accounts for {total_sch
• Rank {self.rank}: Number of indices for this rank ({len(indi
• Error processing sample index {idx}. Error: {e}. Skipping sa
• Calculated batch size <= 0 (seq_len={seq_len}, remaining={re

AI-assisted analysis of huggingface/pytorch-image-models@9a5261e31b (2026-08-27). Data as JSON: /api/errors/68479178dc8e835a. Report an issue: GitHub.