How to create a PyTorch Dataset from a pyspark dataframe (using Databricks)?

Question

I\'m trying to construct a PyTorch Dataset using as raw data a pyspark dataframe (that\'s useful as my training data is stored in columnar format).

I\'m doin