Key error for DatasetTorch

Hi there.

I am trying to follow the tutorial and there seems to be a Key error when using DatasetTorch:

`KeyError Traceback (most recent call last)
Cell In[21], line 8
5 file_list_ = [folder_ + file_ for file_ in os.listdir(folder_)]
7 # Training dataset
----> 8 torch_training_dataset = lanfactory.trainers.DatasetTorch(
9 file_ids=file_list_, batch_size=128
10 )
12 torch_training_dataloader = torch.utils.data.DataLoader(
13 torch_training_dataset,
14 shuffle=True,
(...)
17 pin_memory=True,
18 )
20 # Validation dataset

File ~\anaconda3\envs\LANenv\lib\site-packages\lanfactory\trainers\torch_mlp.py:69, in DatasetTorch.init(self, file_ids, batch_size, label_lower_bound, label_upper_bound, features_key, label_key, out_framework)
66 self.tmp_data = None
68 # get metadata from loading a test file
---> 69 self.__init_file_shape()

File ~\anaconda3\envs\LANenv\lib\site-packages\lanfactory\trainers\torch_mlp.py:118, in DatasetTorch.__init_file_shape(self)
114 def __init_file_shape(self):
115 # Function gets dimensionalities form a test data file
116 init_file = pickle.load(open(self.file_ids[0], "rb"))
117 self.file_shape_dict = {
--> 118 "inputs": init_file[self.features_key].shape,
119 "labels": init_file[self.label_key].shape,
120 }
121 self.batches_per_file = int(self.file_shape_dict["inputs"][0] / self.batch_size)
122 self.input_dim = self.file_shape_dict["inputs"][1]

KeyError: 'data'`

Could this be a bug in the code?

Thanks in advance!

Reproduce the tutorial call to trainers.DatasetTorch and inspect trainers/torch_mlp.py, especially __init_file_shape and the features_key and label_key parameters. Compare the keys in the loaded pickle file with the keys DatasetTorch expects; done means the tutorial dataset initializes without the reported KeyError.