metatensor einops rearrange #6397

import torch
from monai.data import MetaTensor
from einops import rearrange
from monai.transforms import Resize

# create a metatensor in batch mode
x = MetaTensor(torch.randn((1, 1, 481, 404, 3))).to("cuda:0")
x.is_batch = True
x.meta['affine'] = torch.eye(4)[None]

# select the first element in the batch and resize
out = Resize(spatial_size=(1024, 860, -1), align_corners=True, mode="trilinear")(x[0])
print(out.shape)  # works fine

# rearrange using einops
x_ = rearrange(x, "b c h w d -> (b c) h w d")
out = Resize(spatial_size=(1024, 860, -1), align_corners=True, mode="trilinear")(x_)
print(out.shape)  # error

x_ = rearrange(x, "b c h w d -> (b c) h w d")
if isinstance(x_, MetaTensor):
    x_ = x_.as_tensor()
out = Resize(spatial_size=(1024, 860, -1), align_corners=True, mode="trilinear")(x_)
print(out.shape)