Hi,
This might be the dataset module:

import os
import hdf5storage as h5
import torch
from torch.utils.data import DataLoader, Dataset
import torchvision.transforms

class DataSet(Dataset):
	
	def __init__(self, images_folder,transform):
		
		assert os.path.exists(images_folder), 'Images folder does not exist'
		
		self.images_folder=os.path.abspath(images_folder)
		self.images = os.listdir(self.images_folder)
		
		self.transform = transform
		

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

	def __getitem__(self, index):
		
		img_filename = self.images[index]
		img = h5.loadmat(img_filename)['Insert_your_array_name']
		if self.transform:
			img = torch.from_numpy(img)	
		return img


if __name__=='__main__':

	'''
	Folders:
	image				[Training Images]
	image_test			[Testing Images]	
	image_validation	[Validation Images]
	'''

	dataset_train = DataSet(os.path.join(data_path, 'image'), transform)
	dataset_test = DataSet(os.path.join(data_path, 'image_test'), transform)
	dataset_val = DataSet(os.path.join(data_path, 'image_validation'), transform)

	train_loader = DataLoader(dataset_train, batch_size=batch_size, shuffle=True)
    test_loader = DataLoader(dataset_test, batch_size=batch_size, shuffle=True)
    val_loader= DataLoader(dataset_val, batch_size=batch_size, shuffle=True)



Does that solve your problem?