I have implemented a version of this:
MyCollator:
def __init__(self, batch_size, height, width, channels):
        self.buffer = torch.zeros(8, self.batch_size, self.channels, self.height, self.width, dtype=torch.bfloat16)
        self.curr_idx = -1

def __call__(self, batch):
        elem = batch[0]
        elem_type = type(elem)

        if isinstance(elem, torch.Tensor):
            self.curr_idx = (self.curr_idx + 1) % 8
            torch.stack(batch, 0, out=self.buffer[self.curr_idx])
            return self.buffer[self.curr_idx]

collator = MyCollator(self.batch_size, self.in_height, self.in_width, 3)
loader = DataLoader(dataset,
batch_size=self.batch_size,
drop_last=True,
num_workers=self.num_workers,
collate_fn=collator)
This is however returning the same tensor over multiple batches, what could be the issue here?
Also, each worker gets its own copy of the collator, so there should not be any overwriting of the buffer?