Should I create a PyTorch Dataset to train a model off a pyspark dataframe?

Question

I want to train a PyTorch NLP model over training data in columnar format, and I thought to construct a PyTorch Dataset using as raw data a pyspark dataframe (not sure it's the right approach...).

To preprocess text I'm using a tokenizer provided by the transformers library and a tokenizing_UDF function to apply the tokenization.

The Dataset object is then fed to a DataLoader to train a ML model.

What I currently have is this:

import pandas as pd # ideally I'd like to get rid of pandas here
import torch
from torch.utils.data.dataset import Dataset
from transformers import BertTokenizer
from pyspark.sql import types as T
from pyspark.sql import functions as F

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

text = ["This is a test.", "This is not a test."]*100
label = [1, 0]*100

df = sqlContext.createDataFrame(zip(text, label), schema=['text', 'label'])
tokenizing_UDF = udf(lambda t: tokenizer.encode(t),  T.ArrayType(T.LongType())) 
df = df.withColumn("tokenized", tokenizing_UDF(F.col("text"))) # not sure this is the right way
df = df.toPandas() # ugly

class TokenizedDataset(Dataset):
    """needs refactoring..."""
    def __init__(self, df):
        self.data = df
        
    def __getitem__(self, index):
        text = self.data.loc[index].tokenized
        text = torch.LongTensor(text)
        label = self.data.loc[index].label
        return (text, label)

    def __len__(self):
        count = len(self.data)
        return count

dataset = TokenizedDataset(df) # slow...

I currently invoke .toPandas() so that my TokenizedDataset can deal with a pandas dataframe.

Is this a sensible approach? If so, how should I modify the TokenizedDataset code to handle pyspark dataframes directly? If I am off track, should I use https://github.com/uber/petastorm instead?