The MaskModule does not support 3D input

[LVeefkind]

When using the MaskModule on 3D input, the following error is raised:
File "/home/lars/anaconda3/envs/thesis/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl return forward_call(*input, **kwargs) File "/home/lars/anaconda3/envs/thesis/lib/python3.10/site-packages/escnn/nn/modules/sequential_module.py", line 88, in forward x = m(x) File "/home/lars/anaconda3/envs/thesis/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl return forward_call(*input, **kwargs) File "/home/lars/anaconda3/envs/thesis/lib/python3.10/site-packages/escnn/nn/modules/masking_module.py", line 70, in forward assert input.tensor.shape[2:] == self.mask.shape[2:] AssertionError

The two asserted shapes report as follows:
torch.Size([29, 29, 29]) torch.Size([29, 29])

So it seems that the MaskModule does not support 3D input. It would be nice if it also supported 3D input (either dynamically or through a flag/seperate input). Would that be possible?