The number of STL and GPU

[Afreshbird]

I notice that the number of STL is 6 in your code, but the actual number used is 1.
In the forward function of the StageModule class, the return statement is written inside the for loop.
That means no matter what the depth is, always use only one STL.

I try to write the return statement outside the for loop, the consequence was that the GTX 3090 ran out of video memory when I only used one GTX 3090.

But in the IV.A chapter of your paper, “The training platform is with Intel I9-10850K CPU, 64-GB RAM, NVIDIA GeForce GTX 3090 GPU”.
I mean, did you only use one GPU？