f METHOD_NAME(
  self,
  data_rng: spec.RandomState,
  split: str,
  data_dir: str,
  global_batch_size: int,
  cache: Optional[bool] = None,
  repeat_final_dataset: Optional[bool] = None
-> torch.utils.data.DataLoader: