how to use torch dataloader to split data into batches such that they can be aligned with labels well?

Question

I am trying to create batches for torch dataloader on databricks.

My py3 code:

from pyspark.sql.types import StringType, IntegerType, ArrayType

cSchem