I’ve implemented the method you suggested above:
class FitsFolder(DatasetFolder):
EXTENSIONS = ['.fits']

def __init__(self, root, transform=None, target_transform=None,
             loader=None):
    if loader is None:
        loader = self.__fits_loader

    super(FitsFolder, self).__init__(root, loader, self.EXTENSIONS,
                                     transform=transform,
                                     target_transform=target_transform)

@staticmethod
def __fits_loader(filename):
    file = fits.open(filename)
    return (file[1].data)

But when I implement the actual loading process of data:
testset = FitsFolder(root='/Documents/Image_data2')
test_loader = torch.utils.data.DataLoader(testset, batch_size=batch_size,shuffle=False)

And try to plot them, the images look nothing like they are supposed to:

Rather than:

I also don’t know how to perform a transform such as:
transform = transforms.Compose([
                    transforms.Resize((32,32)),
                    transforms.ToTensor(),
                    ])

When using this setup…any help would be greatly appreciated!