Image training error

[ptrblck]

You would just have to use train_data instead of dataset.
Once you run into the error, you can get the image path using train_data.imgs[idx+1].