Yes, you are right about the issues. Here is the updated Dataset function:
class dataset_h5(torch.utils.data.Dataset):
    """
    Reads in a dataset
    """
    def __init__(self, in_file, mode = 'training'):
        super(dataset_h5, self).__init__()
        self.file_path = in_file
        self.dataset_mode = mode
        self.dataset = None

        with h5py.File(self.file_path, 'r') as file:
            self.n_images, self.channels,self.nx, self.ny = file[self.dataset_mode]['inputs'].shape
            self.n_images_check, self.nfeatures = file[self.dataset_mode]['labels'].shape

        if self.n_images != self.n_images_check:
            print('Number of input samples does not match number of label samples!')

        norm_set = h5py.File(in_file, 'r')['normalization']
        self.data_mean = norm_set['mean'][0]
        self.data_std = norm_set['std'][0]
        self.transform =transforms.Normalize((self.data_mean,), (self.data_std,))

    def __getitem__(self, index):
        if self.dataset is None:
            self.dataset = h5py.File(self.file_path, 'r')[self.dataset_mode]
        input = self.transform(torch.tensor(self.dataset['inputs'][index, :, :].astype('float32')))
        labels = torch.tensor(self.dataset['labels'][index, :].astype('float32'))

        return input, labels