I was able to get around this issue by serializing the 2D array my Dataset returns into a string:
data = str(data)
return data
Then in my training loop, I decode the string returned from the DataLoader using the Python ast
library.
data = ast.literal_eval(data)
If your batch size is greater than one, the DataLoader will return a list of strings!