Try this
def _get_samples(self, n):
  samples = self.dataset.data[:n]
  if type(samples) is not torch.Tensor:
    samples = torch.from_numpy(samples)
  return samples

DataLoader.get_samples = _get_samples