Zhangxing Bian

Nvm, I found myself a workaround, pretty simple: add two lines and one helper function: ```python def _get_labels(self, dataset): if self.callback_get_label: return self.callback_get_label(dataset) elif isinstance(dataset, torchvision.datasets.MNIST): return dataset.train_labels.tolist() elif isinstance(dataset,...