Slicing a MetaTensor: TypeError: only integer tensors of a single element can be converted to an index

[jeanmonet]

The goal is to slice the tensor in half. For some reason that I cannot comprehend, it doesn't work on a particular tensor:

>>> type(x)
<class 'monai.data.meta_tensor.MetaTensor'>

>>> x.shape
torch.Size([128, 1024, 8, 64])

>>> x[slice(0, 64, None), ...]
*** TypeError: only integer tensors of a single element can be converted to an index

# OK with [0:16]:
>>> x[slice(0, 16, None), ...].shape
torch.Size([16, 1024, 8, 64])

# NOT OK with [0:32]:
>>> x[slice(0, 32, None), ...]
*** TypeError: only integer tensors of a single element can be converted to an index

>>> x[0:16].shape
torch.Size([16, 1024, 8, 64])

# ARGGGGGHHHH
>>> x[0:32].shape
*** TypeError: only integer tensors of a single element can be converted to an index

>>> x[0:17].shape
*** TypeError: only integer tensors of a single element can be converted to an index

# What if we convert to Numpy? (result: it works fine)
>>> x.detach().numpy()[0:64].shape
(64, 1024, 8, 64)

What could be going on here?

If it is of any use, the error occurs at this specific line: https://github.com/pashtari/factorizer/blob/fa7f03dfe8c390f0993cbc0d161922d2de819a1d/factorizer/factorization/operations.py#L331

To Reproduce

# This works fine, I cannot reproduce in this way
>>> from monai.data import MetaTensor
>>> x = MetaTensor(np.random.rand(128, 1024, 8, 64))
>>> x[slice(0, 64, None), ...].shape
torch.Size([128, 1024, 8, 64])

How is it possible that a specific tensor would cause such error?

Environment

Ubuntu 20.04
python 3.10

monai                     1.1.0              pyhd8ed1ab_0    conda-forge

pytorch                   1.13.1          py3.10_cuda11.7_cudnn8.5.0_0    pytorch
pytorch-cuda              11.7                 h67b0de4_1    pytorch
pytorch-lightning         1.8.1              pyhd8ed1ab_0    conda-forge
pytorch-mutex             1.0                        cuda    pytorch
torchmetrics              0.11.0             pyhd8ed1ab_0    conda-forge
torchvision               0.14.1              py310_cu117    pytorch

[wyli]

thanks for reporting, it's related to the is_batch flag, I'm looking into this issue

>>> x = MetaTensor(np.random.rand(128, 1024, 8, 64))
>>> x.is_batch = True
>>> x[slice(0, 64, None), ...]
2023-01-11 15:45:21,510 - > collate/stack a list of tensors
Traceback (most recent call last):

(it seems the error message I have is different from yours @jeanmonet, could you please share the full traceback if it's convenient for you to do it?)

[jeanmonet]

Thanks for your fast reply. In the meantime I am able to temporarily circumvent the issue by using torch's DataLoader directly:

# Replace:
from monai.data import DataLoader
# With:
from torch.utils.data import DataLoader

could you please share the full traceback

import numpy as np
from monai.data import MetaTensor
x = MetaTensor(np.random.rand(128, 1024, 8, 64))
x.is_batch = True           # https://github.com/Project-MONAI/MONAI/issues/5844
print(x.shape, x.is_batch)
x[slice(0, 64, None), ...].shape

Full traceback

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/monai/data/utils.py:470, in list_data_collate(batch)
    469 else:
--> 470     ret = collate_meta_tensor(data)
    471 return ret

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/monai/data/utils.py:434, in collate_meta_tensor(batch)
    433 collated = default_collate(batch)
--> 434 collated.meta = default_collate([i.meta or TraceKeys.NONE for i in batch])
    435 collated.applied_operations = [i.applied_operations or TraceKeys.NONE for i in batch]

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/torch/utils/data/_utils/collate.py:265, in default_collate(batch)
    205 r"""
    206     Function that takes in a batch of data and puts the elements within the batch
    207     into a tensor with an additional outer dimension - batch size. The exact output type can be
   (...)
    263         >>> default_collate(batch)  # Handle `CustomType` automatically
    264 """
--> 265 return collate(batch, collate_fn_map=default_collate_fn_map)

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/torch/utils/data/_utils/collate.py:128, in collate(batch, collate_fn_map)
    127 try:
--> 128     return elem_type({key: collate([d[key] for d in batch], collate_fn_map=collate_fn_map) for key in elem})
    129 except TypeError:
    130     # The mapping type may not support `__init__(iterable)`.

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/torch/utils/data/_utils/collate.py:128, in <dictcomp>(.0)
    127 try:
--> 128     return elem_type({key: collate([d[key] for d in batch], collate_fn_map=collate_fn_map) for key in elem})
    129 except TypeError:
    130     # The mapping type may not support `__init__(iterable)`.

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/torch/utils/data/_utils/collate.py:120, in collate(batch, collate_fn_map)
    119 if elem_type in collate_fn_map:
--> 120     return collate_fn_map[elem_type](batch, collate_fn_map=collate_fn_map)
    122 for collate_type in collate_fn_map:

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/torch/utils/data/_utils/collate.py:163, in collate_tensor_fn(batch, collate_fn_map)
    162     out = elem.new(storage).resize_(len(batch), *list(elem.size()))
--> 163 return torch.stack(batch, 0, out=out)

RuntimeError: stack expects each tensor to be equal size, but got [4] at entry 0 and [4, 4] at entry 4

The above exception was the direct cause of the following exception:

RuntimeError                              Traceback (most recent call last)
Cell In[75], line 6
      4 x.is_batch = True           # https://github.com/Project-MONAI/MONAI/issues/5844
      5 print(x.shape, x.is_batch)
----> 6 x[slice(0, 64, None), ...].shape

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/monai/data/meta_tensor.py:289, in MetaTensor.__torch_function__(cls, func, types, args, kwargs)
    287 else:
    288     unpack = False
--> 289 ret = MetaTensor.update_meta(ret, func, args, kwargs)
    290 return ret[0] if unpack else ret

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/monai/data/meta_tensor.py:239, in MetaTensor.update_meta(rets, func, args, kwargs)
    237 ret_meta = decollate_batch(args[0], detach=False)[batch_idx]
    238 if isinstance(ret_meta, list):  # e.g. batch[0:2], re-collate
--> 239     ret_meta = list_data_collate(ret_meta)
    240 else:  # e.g. `batch[0]` or `batch[0, 1]`, batch index is an integer
    241     ret_meta.is_batch = False

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/monai/data/utils.py:483, in list_data_collate(batch)
    477         re_str += (
    478             "\n\nMONAI hint: if your transforms intentionally create images of different shapes, creating your "
    479             + "`DataLoader` with `collate_fn=pad_list_data_collate` might solve this problem (check its "
    480             + "documentation)."
    481         )
    482     _ = dev_collate(data)
--> 483     raise RuntimeError(re_str) from re
    484 except TypeError as re:
    485     re_str = str(re)

RuntimeError: stack expects each tensor to be equal size, but got [4] at entry 0 and [4, 4] at entry 4

MONAI hint: if your transforms intentionally create images of different shapes, creating your `DataLoader` with `collate_fn=pad_list_data_collate` might solve this problem (check its documentation).

However, the training Traceback (pointing to this line: https://github.com/pashtari/factorizer/blob/fa7f03dfe8c390f0993cbc0d161922d2de819a1d/factorizer/factorization/operations.py#L332) is:

Full traceback

---------------------------------------------------------------------------
Exception                                 Traceback (most recent call last)
Cell In[92], line 1
----> 1 trainer.fit(
      2     model,
      3     train_dataloaders=train_dataloader,
      4     val_dataloaders=val_dataloader
      5 )

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:582, in Trainer.fit(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
    580     raise TypeError(f"`Trainer.fit()` requires a `LightningModule`, got: {model.__class__.__qualname__}")
    581 self.strategy._lightning_module = model
--> 582 call._call_and_handle_interrupt(
    583     self, self._fit_impl, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path
    584 )

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/pytorch_lightning/trainer/call.py:38, in _call_and_handle_interrupt(trainer, trainer_fn, *args, **kwargs)
     36         return trainer.strategy.launcher.launch(trainer_fn, *args, trainer=trainer, **kwargs)
     37     else:
---> 38         return trainer_fn(*args, **kwargs)
     40 except _TunerExitException:
     41     trainer._call_teardown_hook()

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:624, in Trainer._fit_impl(self, model, train_dataloaders, val_dataloaders, datamodule, ckpt_path)
    617 ckpt_path = ckpt_path or self.resume_from_checkpoint
    618 self._ckpt_path = self._checkpoint_connector._set_ckpt_path(
    619     self.state.fn,
    620     ckpt_path,  # type: ignore[arg-type]
    621     model_provided=True,
    622     model_connected=self.lightning_module is not None,
    623 )
--> 624 self._run(model, ckpt_path=self.ckpt_path)
    626 assert self.state.stopped
    627 self.training = False

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:1061, in Trainer._run(self, model, ckpt_path)
   1057 self._checkpoint_connector.restore_training_state()
   1059 self._checkpoint_connector.resume_end()
-> 1061 results = self._run_stage()
   1063 log.detail(f"{self.__class__.__name__}: trainer tearing down")
   1064 self._teardown()

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:1140, in Trainer._run_stage(self)
   1138 if self.predicting:
   1139     return self._run_predict()
-> 1140 self._run_train()

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:1153, in Trainer._run_train(self)
   1150 self._pre_training_routine()
   1152 with isolate_rng():
-> 1153     self._run_sanity_check()
   1155 # enable train mode
   1156 assert self.model is not None

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:1225, in Trainer._run_sanity_check(self)
   1223 # run eval step
   1224 with torch.no_grad():
-> 1225     val_loop.run()
   1227 self._call_callback_hooks("on_sanity_check_end")
   1229 # reset logger connector

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/pytorch_lightning/loops/loop.py:199, in Loop.run(self, *args, **kwargs)
    197 try:
    198     self.on_advance_start(*args, **kwargs)
--> 199     self.advance(*args, **kwargs)
    200     self.on_advance_end()
    201     self._restarting = False

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/pytorch_lightning/loops/dataloader/evaluation_loop.py:152, in EvaluationLoop.advance(self, *args, **kwargs)
    150 if self.num_dataloaders > 1:
    151     kwargs["dataloader_idx"] = dataloader_idx
--> 152 dl_outputs = self.epoch_loop.run(self._data_fetcher, dl_max_batches, kwargs)
    154 # store batch level output per dataloader
    155 self._outputs.append(dl_outputs)

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/pytorch_lightning/loops/loop.py:199, in Loop.run(self, *args, **kwargs)
    197 try:
    198     self.on_advance_start(*args, **kwargs)
--> 199     self.advance(*args, **kwargs)
    200     self.on_advance_end()
    201     self._restarting = False

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py:137, in EvaluationEpochLoop.advance(self, data_fetcher, dl_max_batches, kwargs)
    134 self.batch_progress.increment_started()
    136 # lightning module methods
--> 137 output = self._evaluation_step(**kwargs)
    138 output = self._evaluation_step_end(output)
    140 self.batch_progress.increment_processed()

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/pytorch_lightning/loops/epoch/evaluation_epoch_loop.py:234, in EvaluationEpochLoop._evaluation_step(self, **kwargs)
    223 """The evaluation step (validation_step or test_step depending on the trainer's state).
    224 
    225 Args:
   (...)
    231     the outputs of the step
    232 """
    233 hook_name = "test_step" if self.trainer.testing else "validation_step"
--> 234 output = self.trainer._call_strategy_hook(hook_name, *kwargs.values())
    236 return output

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/pytorch_lightning/trainer/trainer.py:1443, in Trainer._call_strategy_hook(self, hook_name, *args, **kwargs)
   1440     return
   1442 with self.profiler.profile(f"[Strategy]{self.strategy.__class__.__name__}.{hook_name}"):
-> 1443     output = fn(*args, **kwargs)
   1445 # restore current_fx when nested context
   1446 pl_module._current_fx_name = prev_fx_name

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/pytorch_lightning/strategies/strategy.py:390, in Strategy.validation_step(self, *args, **kwargs)
    388 with self.precision_plugin.val_step_context():
    389     assert isinstance(self.model, ValidationStep)
--> 390     return self.model.validation_step(*args, **kwargs)

File /path/to/script/factorizer/utils/lightning.py:104, in SemanticSegmentation.validation_step(self, batch, batch_idx)
    101 y, id_ = batch["target"], batch["id"]
    103 # inference
--> 104 y_hat = self.predict_step(batch, batch_idx)
    106 # calculate metrics
    107 out = {"id": id_}

File /path/to/script/factorizer/utils/lightning.py:84, in SemanticSegmentation.predict_step(self, batch, batch_idx)
     83 def predict_step(self, batch, batch_idx):
---> 84     y_hat = self.inferer(batch, self)
     85     return y_hat

Cell In[87], line 23, in Inferer.__call__(self, batch, model)
     17 batch = copy.deepcopy(batch)
     19 # infer
     20 
     21 # batch[key_x].is_batch = False   # https://github.com/Project-MONAI/MONAI/issues/5844
---> 23 batch[key_x] = self.inferer(batch[key_x], model)
     24 # postprocess
     25 batch = self.post(batch)

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/monai/inferers/inferer.py:85, in SimpleInferer.__call__(self, inputs, network, *args, **kwargs)
     74 def __call__(self, inputs: torch.Tensor, network: Callable[..., torch.Tensor], *args: Any, **kwargs: Any):
     75     """Unified callable function API of Inferers.
     76 
     77     Args:
   (...)
     83 
     84     """
---> 85     return network(inputs, *args, **kwargs)

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/torch/nn/modules/module.py:1194, in Module._call_impl(self, *input, **kwargs)
   1190 # If we don't have any hooks, we want to skip the rest of the logic in
   1191 # this function, and just call forward.
   1192 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1193         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1194     return forward_call(*input, **kwargs)
   1195 # Do not call functions when jit is used
   1196 full_backward_hooks, non_full_backward_hooks = [], []

File /path/to/script/factorizer/utils/lightning.py:64, in SemanticSegmentation.forward(self, x)
     63 def forward(self, x):
---> 64     out = self.net(x)
     65     if not isinstance(out, torch.Tensor):
     66         out = out[0]

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/torch/nn/modules/module.py:1194, in Module._call_impl(self, *input, **kwargs)
   1190 # If we don't have any hooks, we want to skip the rest of the logic in
   1191 # this function, and just call forward.
   1192 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1193         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1194     return forward_call(*input, **kwargs)
   1195 # Do not call functions when jit is used
   1196 full_backward_hooks, non_full_backward_hooks = [], []

File /path/to/script/factorizer/unet.py:502, in UNet.forward(self, x)
    501 def forward(self, x):
--> 502     y = self.forward_features(x)
    504     out = []
    505     for j, head in enumerate(self.heads):

File /path/to/script/factorizer/unet.py:497, in UNet.forward_features(self, x)
    495 def forward_features(self, x):
    496     out = self.stem(x)
--> 497     out = self.encoder(out)
    498     out = self.decoder(out)
    499     return out

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/torch/nn/modules/module.py:1194, in Module._call_impl(self, *input, **kwargs)
   1190 # If we don't have any hooks, we want to skip the rest of the logic in
   1191 # this function, and just call forward.
   1192 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1193         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1194     return forward_call(*input, **kwargs)
   1195 # Do not call functions when jit is used
   1196 full_backward_hooks, non_full_backward_hooks = [], []

File /path/to/script/factorizer/unet.py:332, in UNetEncoder.forward(self, x)
    330 def forward(self, x):
    331     print("Unet encoder block 0")
--> 332     out = [self.blocks[0](x)]
    333     for i, block in enumerate(self.blocks[1:], start=1):
    334         print("Unet encoder block", i)

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/torch/nn/modules/module.py:1194, in Module._call_impl(self, *input, **kwargs)
   1190 # If we don't have any hooks, we want to skip the rest of the logic in
   1191 # this function, and just call forward.
   1192 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1193         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1194     return forward_call(*input, **kwargs)
   1195 # Do not call functions when jit is used
   1196 full_backward_hooks, non_full_backward_hooks = [], []

File /path/to/script/factorizer/unet.py:283, in UNetEncoderBlock.forward(self, x)
    282 def forward(self, x):
--> 283     x = self.blocks(x)
    284     return x

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/torch/nn/modules/module.py:1194, in Module._call_impl(self, *input, **kwargs)
   1190 # If we don't have any hooks, we want to skip the rest of the logic in
   1191 # this function, and just call forward.
   1192 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1193         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1194     return forward_call(*input, **kwargs)
   1195 # Do not call functions when jit is used
   1196 full_backward_hooks, non_full_backward_hooks = [], []

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/torch/nn/modules/container.py:204, in Sequential.forward(self, input)
    202 def forward(self, input):
    203     for module in self:
--> 204         input = module(input)
    205     return input

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/torch/nn/modules/module.py:1194, in Module._call_impl(self, *input, **kwargs)
   1190 # If we don't have any hooks, we want to skip the rest of the logic in
   1191 # this function, and just call forward.
   1192 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1193         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1194     return forward_call(*input, **kwargs)
   1195 # Do not call functions when jit is used
   1196 full_backward_hooks, non_full_backward_hooks = [], []

File /path/to/script/factorizer/factorizer.py:174, in FactorizerBlock.forward(self, x)
    171 out = self.pos_embed(out) if hasattr(self, "pos_embed") else out
    173 for block in self.blocks.values():
--> 174     out = block(out)
    176 return out

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/torch/nn/modules/module.py:1194, in Module._call_impl(self, *input, **kwargs)
   1190 # If we don't have any hooks, we want to skip the rest of the logic in
   1191 # this function, and just call forward.
   1192 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1193         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1194     return forward_call(*input, **kwargs)
   1195 # Do not call functions when jit is used
   1196 full_backward_hooks, non_full_backward_hooks = [], []

File /path/to/script/factorizer/factorizer.py:59, in Residual.forward(self, x, **kwargs)
     58 def forward(self, x, **kwargs):
---> 59     return self.fn(x, **kwargs) + x

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/torch/nn/modules/module.py:1194, in Module._call_impl(self, *input, **kwargs)
   1190 # If we don't have any hooks, we want to skip the rest of the logic in
   1191 # this function, and just call forward.
   1192 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1193         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1194     return forward_call(*input, **kwargs)
   1195 # Do not call functions when jit is used
   1196 full_backward_hooks, non_full_backward_hooks = [], []

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/torch/nn/modules/container.py:204, in Sequential.forward(self, input)
    202 def forward(self, input):
    203     for module in self:
--> 204         input = module(input)
    205     return input

File /opt/miniconda3/envs/py10/lib/python3.10/site-packages/torch/nn/modules/module.py:1194, in Module._call_impl(self, *input, **kwargs)
   1190 # If we don't have any hooks, we want to skip the rest of the logic in
   1191 # this function, and just call forward.
   1192 if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1193         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1194     return forward_call(*input, **kwargs)
   1195 # Do not call functions when jit is used
   1196 full_backward_hooks, non_full_backward_hooks = [], []

File /path/to/script/factorizer/factorizer.py:120, in FactorizerSubblock.forward(self, x)
    117 out = out.reshape(shape)
    119 # detensorize (unfold/fold)
--> 120 out = self.tensorize.inverse_forward(out)
    122 # output linear projection
    123 out = self.out_proj(out)

File /path/to/script/factorizer/factorization/operations.py:340, in SWMatricize.inverse_forward(self, x)
    338     print("slice:", type(slc), slc, j, b, num_shifts)
    339     print("x shape", x.shape)
--> 340     raise Exception(str(exp)) from None
    341 else:
    342     if not self.FLAG:

Exception: only integer tensors of a single element can be converted to an index
---------------------------------------------------------------------------

[wyli]

thanks, I still couldn't get the TypeError only integer tensors of a single element can be converted to an index... But in my understanding, with the factorization library, the minibatch images are reshaped into vectorized patches and in this case using the metadata to track the spatial relationship is not defined/implemented. So I'd recommend using x.as_tensor() to convert metatensor x into a regular pytorch tensor for that operation.

you can still track the meta info using x.meta:

import factorizer as ft
import torch

from monai.data import MetaTensor
from monai.data.utils import collate_meta_tensor

x = [MetaTensor(torch.rand(1024, 8, 64)) for _ in range(128)]
x = collate_meta_tensor(x)
m = x.meta.copy()
sw_matricize = ft.SWMatricize((128, 1024, 8, 64), head_dim=8, patch_size=8)
x = x.as_tensor()
y = sw_matricize(x)  # (262144, 8, 64)
z = sw_matricize.inverse_forward(y)
z.meta = m
print(torch.equal(x, z))  # True

I'll make a PR to improve the error message about the usage.

[jeanmonet]

Thank you for investigating and for the suggestion, it is well noted and very helpful.