Doing the sampling of batch indexes inside lightning

[pamparana34]

I am trying to port my very old pytorch code to lightning and in my training loop, I have something as follows:

batch_order = np.arange(data.x_train.shape[0])
batch_index = np.random.choice(batch_order, batch_size, p=seq_sample_probs).tolist()
batch = torch.tensor(data.x_train[batch_index], dtype=torch_dtype, device=torch_device, requires_grad=False)

# then call the forward on this batch
model.encoder.forward(batch)

I was wondering how I can incorporate this batch index selection in the lightning code. In my code, I have the usual:

def train_dataloader(self):
        return DataLoader(self.train_dataset, batch_size=self.batch_size)

But I do not know where I can inject my sampling code inside all this.

[rohitgr7]

I think that's the job of the PyTorch sampler.