LabelMap not tranformed in tio.Subject when label shape is 1xHxWx1

🐛 Describe the bug

Situation: In a dataloader a tio composition of intensity and spatial augmentation is performed.
When using 3D samples the tio.ScalarImage and tio.LabelMap are transformed fine, when using 2D samples the tio.LabelMap is not spatially transformed. This only happens when within the tio.Compose a intensity transform is applied along with spatial transform (either prior or after spatial transformation).

import torchio as tio

if self.yield_2d:            
    self.spatial_transform = tio.OneOf({
        tio.transforms.RandomAffine(.05): 0.8,
        tio.transforms.RandomElasticDeformation(num_control_points=7, max_displacement=(7.5,7.5,0)): 0.2,
        },
        p=0.75,
    )
else:
    self.spatial_transform = tio.OneOf({
        tio.transforms.RandomAffine(.05): 0.8,
        tio.transforms.RandomElasticDeformation(num_control_points=7, max_displacement=7.5): 0.2,
        },
        p=0.75,
    )

self.transform = tio.Compose([
    self.intensity_transform,
    self.spatial_transform
])

# 2D image are dim=1xHxWx1 after unsqueezing
# 3D image are dim=1xDxHxW after unsqueezing

if yield_2d:
    image = image.unsqueeze(0).unsqueeze(-1)
    label = label.unsqueeze(0).unsqueeze(-1)
else:
    image = image.unsqueeze(0)
    label = label.unsqueeze(0)

subject = tio.Subject(
    image=tio.ScalarImage(tensor=image),  
    label=tio.LabelMap(tensor=label)
)
subject = self.spatial_transform(subject)

3D images are transformed fine:

2D images are not:

Versions

Collecting environment information...
PyTorch version: 1.9.1+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 16.04.7 LTS (x86_64)
GCC version: (Ubuntu 5.4.0-6ubuntu1~16.04.12) 5.4.0 20160609
Clang version: 3.8.0-2ubuntu4 (tags/RELEASE_380/final)
CMake version: version 3.5.1
Libc version: glibc-2.23

Python version: 3.9.5 (default, Jun 4 2021, 12:28:51) [GCC 7.5.0] (64-bit runtime)
Python platform: Linux-4.4.0-210-generic-x86_64-with-glibc2.23
Is CUDA available: True
CUDA runtime version: 7.5.17
GPU models and configuration:
GPU 0: Tesla T4
GPU 1: Tesla T4
GPU 2: GeForce RTX 2080 Ti
GPU 3: Quadro RTX 8000
GPU 4: GeForce RTX 2080 Ti
GPU 5: GeForce RTX 2080 Ti
GPU 6: GeForce RTX 2080 Ti

Nvidia driver version: 455.45.01
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.0
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn.so.8.1.0
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.1.0
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.1.0
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.1.0
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.1.0
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.1.0
/usr/local/cuda-11.2/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.1.0
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip3] numpy==1.21.1
[pip3] torch==1.9.1
[pip3] torchio==0.18.63
[pip3] torchvision==0.10.1
[conda] Could not collect

Reproduce the example using tio.Subject with ScalarImage and LabelMap tensors shaped 1xHxWx1, comparing Compose with intensity and spatial transforms against the 3D case. Trace how tio.Compose, OneOf, RandomAffine, and RandomElasticDeformation handle the LabelMap, then verify that the 2D label receives the same spatial transformation as the image.