Training Stage 2

[Rawal Khirodkar]

I cannot fit the default training configuration of stage 2 into 4 x 3090 GPUs.

I am using a single node.
BATCH_SIZE_3D: 4, BATCH_SIZE_2D: 3, BATCH_SIZE_IMG: 7
And I am getting out of memory error for Stage 2.

The paper mentions that the stage 2 training was done with 16 x 1080 GPUs.
This should fit into the 3090s as well.

I would greatly appreciate your help.