Hey Yin, I agree, having both TR and Pytorch is not a good idea, not only in the same script but also in the same python module.
I don’t have the exact code anywhere near as it was a while ago since I used tfrecords, but databricks seems to have some good documentation. See it below:
from pyspark.sql.types import *

schema = StructType([StructField('image/class/label', IntegerType(), True),
                     StructField('image/width', IntegerType(), True),
                     StructField('image/height', IntegerType(), True),
                     StructField('image/format', StringType(), True),
                     StructField('image/encoded', BinaryType(), True)])

df = spark.read.format("tfrecords").schema(schema).load(input_local_dir+'/flowers_train*.tfrecord')
df = df.limit(3200)
display(df)

In general, the idea is to read the tfrecords to a pyspark dataframe (do same preprocessing on it) and then save it as avro with df.write.format("avro").save(/tmp/loc).
Link to the databricks notebook on tfrecords is here 12.
Also, happy to see people using pytorch and spark, big fan of the combination!