Logging information from Dataset?

I’d like to log various information about each dataset “record” consumed during the training loop. This includes the idx that was passed from the DataLoader, plus various detailed information such as the exact augmentations that were applied, how long it took to produce the record, etc.

My current solution is to return this information from the Dataset by combining it with the label as a 2-tuple, and then I unpack that in the training loop to separate the label from this extra information. So instead of the dataset returning img, label, it returns img, (label, extra_information).

See last line in block of code below:

class CustomDataset(Dataset):
    def __init__(self, csv_with_image_label_pairs):
        self.csv_with_image_label_pairs = csv_with_image_label_pairs

        df = pd.read_csv(self.csv_with_image_label_pairs)

    def __len__(self):
        return len(self.df)

    def produce_img(img_filepath):
        # lots of complex code
        pass

    def __getitem__(self, idx):
        record = self.df.iloc[idx]

        img = self.produce_image(record["img_filepath"])
        label = record["label"]
        extra_information = (idx, augmentation_params, other_stuff)

        return img, (label, extra_information)

Is this a good way of doing this? Any other approaches to consider? I’m an “advanced beginner” with Python and a beginner with Pytorch, btw.

Thanks!