SwinUNETR pretraining crashes with larger --feature_size

[hw-ju]

Hi!
I'm trying to run single-GPU SwinUNETR pretraining using dataset TCIA Covid 19 only following the tutorial here with different --feature_size by MONAI Core v0.9.1 container (currently the cluster I'm using doesn't have required driver installed to run v1.2.0 container).

It works with the default --feature_size by the command below:

python main.py \
--roi_x=128 --roi_y=128 --roi_z=128 \
--lrdecay --lr=6e-6 \
--batch_size=1 \
--epochs=3 --num_steps=6 --eval_num=2 \
--noamp

However, when I setting the --feature_size = 96 by the command below:

python main.py \
--roi_x=128 --roi_y=128 --roi_z=128 \
--lrdecay --lr=6e-6 \
--batch_size=1 \
--epochs=3 --num_steps=6 --eval_num=2 \
--noamp \
--feature_size=96

an error occurs:

[30595] single-GPU training
[0] current gpu: 0
[0] Writing Tensorboard logs to /mnt/runs/
[0] Dataset 2 Covid 19: number of data: 722
[0] Dataset all training: number of data: 722
[0] Dataset all validation: number of data: 49
[0] Using generic dataset
[0] Total parameters count: 42628663
[0] while loop: new epoch, global_step 0 -------------------------
Traceback (most recent call last):
  File "main.py", line 391, in <module>
    main()
  File "main.py", line 364, in main
    global_step, loss, best_val = train(args, global_step, train_loader, best_val, scaler, count_epoch)
  File "main.py", line 51, in train
    rot1_p, contrastive1_p, rec_x1 = model(x1_augment)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/home/hju/monai_uf_tutorials/pretrain_modify/models/ssl_head.py", line 85, in forward
    x_rot = self.rotation_head(x_rot)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/opt/conda/lib/python3.8/site-packages/torch/nn/modules/linear.py", line 114, in forward
    return F.linear(input, self.weight, self.bias)
RuntimeError: mat1 and mat2 shapes cannot be multiplied (2x1536 and 768x4)

Any ideas?

[KumoLiu]

Hi @hw-ju, you should also modify the dim here.
https://github.com/Project-MONAI/research-contributions/blob/159e9232269b50b1bdd5a5f872c0143cf738ec0b/SwinUNETR/Pretrain/models/ssl_head.py#L20

Hope it can help you, thanks!

[hw-ju]

@KumoLiu Thank you so much!!