Hey @ptrblck, thanks for the reply.
The issue is technical as I am using a stateful dataset where each instance has information stored and processed at each batch. Hence when the same instance is found twice in a mini-batch I have all sorts of errors that arise that I’d rather solve via the sampler rather than modifying my current implementation.
I’ve implemented a sampler that seems to work for now but is probably not that optimal. Any thoughts on it?
from torch.utils.data import WeightedRandomSampler
from typing import Sequence
import torch

class CustomWeightedSampler(WeightedRandomSampler):
    def __init__(self, weights: Sequence[float], num_samples: int, bs: int,
                 generator=None) -> None:
        if not isinstance(num_samples, int) or isinstance(num_samples, bool) or \
                num_samples <= 0:
            raise ValueError("num_samples should be a positive integer "
                             "value, but got num_samples={}".format(num_samples))
        if not isinstance(bs, int) or isinstance(bs, bool) or \
                num_samples <= 0:
            raise ValueError("bs should be a positive integer "
                             "value, but got bs={}".format(num_samples))
        if bs > num_samples:
            raise ValueError("bs should be smaller than num_samples "
                             "but got bs={} and num_samples={}".format(bs, num_samples))
        self.weights = weights
        self.num_samples = num_samples
        self.bs = bs
        self.generator=generator


    def __iter__(self):
        rand_tensor = torch.multinomial(self.weights, self.bs, False, generator=self.generator)
        for _ in range((self.num_samples - self.bs) // self.bs):
            rand_tensor = torch.cat([rand_tensor,
                                      torch.multinomial(self.weights, self.bs, False, generator=self.generator)])
        if self.num_samples % self.bs != 0:
            rand_tensor = torch.cat([rand_tensor,
                                      torch.multinomial(self.weights,  self.num_samples % self.bs, False, generator=self.generator)])
        return iter(rand_tensor.tolist())

    def __len__(self):
        return super(CustomWeightedSampler, self).__len__()