This is the dataset class:

class dataset_class(torch.utils.data.Dataset):
    ''' General dataset class
        Args:
                X: input -> torch.tensor
                Y: targets -> torch.tensor
    '''

    def __init__(self,X: torch.tensor, Y: torch.tensor):
        super(dataset_class, self).__init__()
        assert X is not None , "Invalid None type"
        assert Y is not None, "Invalid None type"
        self.X = X
        self.Y = Y

    def __len__(self):
        return self.X.size(0)

    def __getitem__(self,idx):
      
        return self.X[idx],self.Y[idx]

where  self.X and self.Y are constant for all the runs, i.e its order is not shuffled unless the dataloader does it. This is the main loop:
for b,(x,y) in enumerate(train_loader):

     x,y = x.to(cg.device),y.to(cg.device)
                   
      ## loss
      loss = model.loss(x,y)
      optimizer.zero_grad() # just in case
      loss.backward()
      optimizer.step()
              
      # Performance metrics                      
      if ((ep+1)%self.validate_each) == 0:
           model.performance_metrics(x,y,dataset='train')
 
      with torch.no_grad():

          if ((ep+1)%validate_each) == 0:
                    for b,(x,y) in enumerate(test_loader):
                        x,y = x.to(cg.device),y.to(cg.device)
                        model.performance_metrics(x,y,dataset='test')