question: Doing the sampling of batch indexes inside lightning
context: 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.
answer: I think that's the job of the PyTorch sampler.