Hi,
I want to check 2 things.

Is temp(50) correct? I think it missed one argument according to the definition: temp(data, batch_size).
Does the temp function work as you expected? Is it OK to ignore i in sampled = torch.randperm(len(data))[0:batch_size]?

For me, the snippet below looks natural.
def temp(data, batch_size):
    indices = troch.randperm(len(data))
    for i in range(0, data.shape[0] // batch_size):
        yield i, indices[i * batch_size: (i + 1) * batch_size]