UNETR error because of input image size

[YingdiZhang]

Thanks for the great job. I am using UNETR for a 3D segmentation task. When set img_size=(64,64,64), an error ocurs and I have no idea:
Traceback (most recent call last): File "/data/zyd/beifen_xinsuoqu_/MedSeg/MedSeg/train/Task961TrainerDDP.py", line 1234, in <module> trainer.train(max_epochs=max_epochs, File "/data/zyd/beifen_xinsuoqu_/MedSeg/MedSeg/train/Task961TrainerDDP.py", line 127, in train outputs = self.model(batch_data["image"].to(device)) File "/home/zyd/anaconda3/envs/monai1.1/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl return forward_call(*input, **kwargs) File "/home/zyd/anaconda3/envs/monai1.1/lib/python3.9/site-packages/torch/nn/parallel/distributed.py", line 1040, in forward output = self._run_ddp_forward(*inputs, **kwargs) File "/home/zyd/anaconda3/envs/monai1.1/lib/python3.9/site-packages/torch/nn/parallel/distributed.py", line 1000, in _run_ddp_forward return module_to_run(*inputs[0], **kwargs[0]) File "/home/zyd/anaconda3/envs/monai1.1/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl return forward_call(*input, **kwargs) File "/home/zyd/anaconda3/envs/monai1.1/lib/python3.9/site-packages/monai/networks/nets/unetr.py", line 195, in forward x, hidden_states_out = self.vit(x_in) File "/home/zyd/anaconda3/envs/monai1.1/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl return forward_call(*input, **kwargs) File "/home/zyd/anaconda3/envs/monai1.1/lib/python3.9/site-packages/monai/networks/nets/vit.py", line 110, in forward x = self.patch_embedding(x) File "/home/zyd/anaconda3/envs/monai1.1/lib/python3.9/site-packages/torch/nn/modules/module.py", line 1194, in _call_impl return forward_call(*input, **kwargs) File "/home/zyd/anaconda3/envs/monai1.1/lib/python3.9/site-packages/monai/networks/blocks/patchembedding.py", line 117, in forward embeddings = x + self.position_embeddings File "/home/zyd/anaconda3/envs/monai1.1/lib/python3.9/site-packages/monai/data/meta_tensor.py", line 268, in __torch_function__ ret = super().__torch_function__(func, types, args, kwargs) File "/home/zyd/anaconda3/envs/monai1.1/lib/python3.9/site-packages/torch/_tensor.py", line 1279, in __torch_function__ ret = func(*args, **kwargs) RuntimeError: The size of tensor a (64) must match the size of tensor b (576) at non-singleton dimension 1 ERROR:torch.distributed.elastic.multiprocessing.api:failed (exitcode: 1) local_rank: 0 (pid: 25347) of binary: /home/zyd/anaconda3/envs/monai1.1/bin/python Traceback (most recent call last): File "/home/zyd/anaconda3/envs/monai1.1/lib/python3.9/runpy.py", line 197, in _run_module_as_main return _run_code(code, main_globals, None, File "/home/zyd/anaconda3/envs/monai1.1/lib/python3.9/runpy.py", line 87, in _run_code exec(code, run_globals) File "/home/zyd/anaconda3/envs/monai1.1/lib/python3.9/site-packages/torch/distributed/launch.py", line 195, in <module> main() File "/home/zyd/anaconda3/envs/monai1.1/lib/python3.9/site-packages/torch/distributed/launch.py", line 191, in main launch(args) File "/home/zyd/anaconda3/envs/monai1.1/lib/python3.9/site-packages/torch/distributed/launch.py", line 176, in launch run(args) File "/home/zyd/anaconda3/envs/monai1.1/lib/python3.9/site-packages/torch/distributed/run.py", line 753, in run elastic_launch( File "/home/zyd/anaconda3/envs/monai1.1/lib/python3.9/site-packages/torch/distributed/launcher/api.py", line 132, in __call__ return launch_agent(self._config, self._entrypoint, list(args)) File "/home/zyd/anaconda3/envs/monai1.1/lib/python3.9/site-packages/torch/distributed/launcher/api.py", line 246, in launch_agent raise ChildFailedError(

[tangy5]

This log message: RuntimeError: The size of tensor a (64) must match the size of tensor b (576) , means your input data are somehow mismatched the network's config, not sure where the 576 from, but your could check whether the network configs are correctly changed.
If you are not sure, you could start with this tutorial:
https://github.com/Project-MONAI/tutorials/blob/main/3d_segmentation/unetr_btcv_segmentation_3d.ipynb