You could write a custom Dataset as explained here and apply the transformation in the __getitem__ method.
Since you’ve already created the TensorDataset, you could also pass it to your custom Dataset, add transforms.ToPILImage() as the first transformation, and apply it in the __getitem__ method.