Run on CIFAR-10 failed

[Qing21]

Hi,
I tried to run the capsule network on cifar-10 dataset and modified the config: input height&width ->32 and channels->3. But failed with the error:

Traceback (most recent call last):
File "/ProgramData/Pycharm/capsule_network_dynamic_routing-master/train.py", line 232, in
train()
File "/ProgramData/Pycharm/capsule_network_dynamic_routing-master/train.py", line 47, in train
loss = capsule_loss(data, target, v_length, x_reconst)
File "/Users/opt/anaconda3/lib/python3.7/site-packages/torch/nn/modules/module.py", line 550, in call
result = self.forward(*input, **kwargs)
File "/Users/ProgramData/Pycharm/capsule_network_dynamic_routing-master/capsnet_v2.py", line 168, in forward
assert torch.numel(images) == torch.numel(reconstructions)
AssertionError

Could you please give me some suggestions to successfully run the code? Is there any other parameter or function that need to be changed in order to test on CIFAR-10 dataset?