Hey Yin, spark to torch dataloader does require some custom work but is fairly easy to build. How I do it is I use torch.utils.data.IterableDataset. Firstly I load all the avro/parquet (as you are working with spark) to a DataReader object which is a generator (where I do some of my custom processing on each record). Then I build my own collate_fn function to load the data from a generator to a DataLoader object and return a data loader. Below is a brief sketch of how I do it:
from typing import Dict, Any, Optional, List, Callable

import torch
from fastavro import reader
from functools import partial
from torch.utils.data import DataLoader

class AvroDataReader(torch.utils.data.IterableDataset):

    def __init__(self, data_location: str, custom_fn: Callable):
        self.data_location = data_location
        self.custom_fn = custom_fn

    def __iter__(self):
        avro_files = get_files(self.data_location, '*.avro')
        for filename in avro_files:
            with file_reader(filename) as local_file:
                with smart_open(local_file, 'rb') as i:
                    for record in reader(i):
                        sample = self.custom_fn(record)
                        if sample is not None:
                            yield sample


def _collate_examples(batch_of_data):
    return SomeType(
        token_ids=torch.stack([b.token_ids for b in batch_of_data]),
        label_idx=torch.stack([b.label_idx for b in batch_of_data]),
    )


def get_dataloader(dataset, batch_size: int,
                   collate_fn: Callable):
    return DataLoader(dataset, batch_size=batch_size, collate_fn=collate_fn)


def get_torch_dataloader(data_location: str, batch_size: int, label_vocab: Dict[str, int]):
    dataset = AvroDataReader(data_location=data_location,
                             custom_fn=partial(custom, label_vocab=label_vocab))
    return get_dataloader(dataset=dataset, batch_size=batch_size,
                          collate_fn=_collate_labelled_examples)

You need to obviously adapt your code to the example I provided, but should give you a good sketch of how it should work and should not throw any memory issues.