Rotate function from torchvision doesn't work with AugmentedTensors when setting "fill" parameter

[Dee61298]

Minimal reproducible example :

import torch
import torchvision.transforms.functional as F
import matplotlib.pyplot as plt
from aloscene.tensors import AugmentedTensor


x = torch.rand((1, 3, 100, 200))

y = AugmentedTensor(x, names=("B", "C", "H", "W"))

y = y.rename_(None, auto_restore_names=True)

y = F.rotate(y, 45, fill=0)


y_to_plot = torch.transpose(y[0], 0, 2)
plt.imshow(y_to_plot), plt.show()

Using F.rotate with AugmentedTensor when not passing the "fill" parameter works.
Using F.rotate with the "fill" parameter and with a regular tensor (not AugmentedTensor) works.
Using F.rotate with the "fill" parameter and with AugmentedTensor gives the following stacktrace :

Traceback (most recent call last):
  File "testing.py", line 13, in <module>
    y = F.rotate(y[0], 45, fill=0)
  File "/home/visualbehavior/anaconda3/envs/test/lib/python3.8/site-packages/torchvision/transforms/functional.py", line 1019, in rotate
    return F_t.rotate(img, matrix=matrix, interpolation=interpolation.value, expand=expand, fill=fill)
  File "/home/visualbehavior/anaconda3/envs/test/lib/python3.8/site-packages/torchvision/transforms/functional_tensor.py", line 754, in rotate
    return _apply_grid_transform(img, grid, interpolation, fill=fill)
  File "/home/visualbehavior/anaconda3/envs/test/lib/python3.8/site-packages/torchvision/transforms/functional_tensor.py", line 665, in _apply_grid_transform
    img[mask] = fill_img[mask]
  File "/home/visualbehavior/aloception/aloscene/tensors/augmented_tensor.py", line 302, in __getitem__
    tensor = torch.Tensor.__getitem__(self.rename(None), idx).reset_names()
  File "/home/visualbehavior/aloception/aloscene/tensors/augmented_tensor.py", line 558, in __torch_function__
    tensor = super().__torch_function__(func, types, args, kwargs)
  File "/home/visualbehavior/anaconda3/envs/test/lib/python3.8/site-packages/torch/_tensor.py", line 1051, in __torch_function__
    ret = func(*args, **kwargs)
IndexError: The shape of the mask [3, 150, 200] at index 0 does not match the shape of the indexed tensor [1, 3, 150, 200] at index 0

[ragier]

New minimal example :

import torch
from aloscene.tensors import AugmentedTensor

x = torch.rand((1, 3, 100, 200))
y = AugmentedTensor(x, names=("B", "C", "H", "W"))
mask = y[:, 0, :, :] # 1 * 1 * 100 * 200
mask = mask.expand_as(y)
mask = mask > 0.5
y[mask] = y[mask]