Confusion about order of input dimensions in SwinUNETR

[HSILA]

Consider the SwinUNETR model, the pertaining section. In utils/ops.py, one can have two interpretations of the order of dimensions of the input data after transformations.

1- C,H,W,Z: as in patch_rand_drop method.

2- C,Z,H,W: In the rot_rand method, the rotation (x.rot90(1, (2, 3))) is on axis (2,3), and it is mentioned in the paper that rotations are along the z-axis.

Based on the paper, the second one should be correct, but patch_rand_drop consistently uses the first notation, I'm a little bit confused.