Error when trying to train PoseC3D model on custom dataset

[Zeev1988]

Hi,
I created a custom dataset as defined in prepare skeleton dataset

I run the following code

from mmengine import Config
from mmengine.runner import Runner

config_file = r'./data/custom_config.py'
cfg = Config.fromfile(config_file)
cfg.work_dir = './results'
runner = Runner.from_cfg(cfg)
runner.train()

and get the following error:

05/14 21:15:52 - mmengine - INFO - 2602 videos remain after valid thresholding
05/14 21:16:33 - mmengine - INFO - 2602 videos remain after valid thresholding
05/14 21:16:33 - mmengine - WARNING - "FileClient" will be deprecated in future. Please use io functions in https://mmengine.readthedocs.io/en/latest/api/fileio.html#file-io
05/14 21:16:33 - mmengine - INFO - Checkpoints will be saved to /home/ec2-user/SageMaker/pose3d/results.

AssertionError Traceback (most recent call last)
Cell In[65], line 31
22 runner = Runner.from_cfg(cfg)
23 # Start fine-tuning
24 # train_model(
25 # model,
(...)
29 # validate=True
30 # )
---> 31 runner.train()

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/mmengine/runner/runner.py:1777, in Runner.train(self)
1773 # Maybe compile the model according to options in self.cfg.compile
1774 # This must be called AFTER model has been wrapped.
1775 self._maybe_compile('train_step')
-> 1777 model = self.train_loop.run() # type: ignore
1778 self.call_hook('after_run')
1779 return model

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/mmengine/runner/loops.py:96, in EpochBasedTrainLoop.run(self)
93 self.runner.call_hook('before_train')
95 while self._epoch < self._max_epochs and not self.stop_training:
---> 96 self.run_epoch()
98 self._decide_current_val_interval()
99 if (self.runner.val_loop is not None
100 and self._epoch >= self.val_begin
101 and (self._epoch % self.val_interval == 0
102 or self._epoch == self._max_epochs)):

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/mmengine/runner/loops.py:112, in EpochBasedTrainLoop.run_epoch(self)
110 self.runner.call_hook('before_train_epoch')
111 self.runner.model.train()
--> 112 for idx, data_batch in enumerate(self.dataloader):
113 self.run_iter(idx, data_batch)
115 self.runner.call_hook('after_train_epoch')

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/utils/data/dataloader.py:630, in _BaseDataLoaderIter.next(self)
627 if self._sampler_iter is None:
628 # TODO(pytorch/pytorch#76750)
629 self._reset() # type: ignore[call-arg]
--> 630 data = self._next_data()
631 self._num_yielded += 1
632 if self._dataset_kind == _DatasetKind.Iterable and
633 self._IterableDataset_len_called is not None and
634 self._num_yielded > self._IterableDataset_len_called:

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/utils/data/dataloader.py:1345, in _MultiProcessingDataLoaderIter._next_data(self)
1343 else:
1344 del self._task_info[idx]
-> 1345 return self._process_data(data)

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/utils/data/dataloader.py:1371, in _MultiProcessingDataLoaderIter._process_data(self, data)
1369 self._try_put_index()
1370 if isinstance(data, ExceptionWrapper):
-> 1371 data.reraise()
1372 return data

File ~/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/_utils.py:694, in ExceptionWrapper.reraise(self)
690 except TypeError:
691 # If the exception takes multiple arguments, don't try to
692 # instantiate since we don't know how to
693 raise RuntimeError(msg) from None
--> 694 raise exception

AssertionError: Caught AssertionError in DataLoader worker process 0.
Original Traceback (most recent call last):
File "/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/utils/data/_utils/worker.py", line 308, in _worker_loop
data = fetcher.fetch(index)
File "/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py", line 51, in fetch
data = [self.dataset[idx] for idx in possibly_batched_index]
File "/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/torch/utils/data/_utils/fetch.py", line 51, in
data = [self.dataset[idx] for idx in possibly_batched_index]
File "/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/mmengine/dataset/dataset_wrapper.py", line 290, in getitem
return self.dataset[sample_idx]
File "/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/mmengine/dataset/base_dataset.py", line 410, in getitem
data = self.prepare_data(idx)
File "/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/mmengine/dataset/base_dataset.py", line 793, in prepare_data
return self.pipeline(data_info)
File "/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/mmengine/dataset/base_dataset.py", line 60, in call
data = t(data)
File "/home/ec2-user/anaconda3/envs/pytorch_p310/lib/python3.10/site-packages/mmcv/transforms/base.py", line 12, in call
return self.transform(results)
File "/home/ec2-user/SageMaker/pose3d/mmaction2/mmaction/datasets/transforms/pose_transforms.py", line 63, in transform
assert k in results
AssertionError

Any idea what could cause this error?
Does the dataset have to be homogeneous regarding the number of frames and video dimensions?