Double Training Step to Swap Inputs

[pederismo]

Hi everyone!

I would like to do something in my training step but I'm not really sure how to do it, or whether it is even possible to do it.
At each training step my batch contains two volumetric images, and they are used in two ways:

• Together (to interpolate the two embeddings);
• Singularly, where volume_1 gets used in algorithm_1 and volume 2 in algorithm_2; they contribute differently to the final loss.

Now if my explanation was clear enough, you should understand that I'm basically halving my dataset: for every batch of two volumetric images I need one training step where I'm using volume_1 in algorithm_1 and volume_2 in algorithm_2, and another training step where I'm using volume_1 in algorithm_2 and volume_2 in algorithm_1.

But I get the batch WITHIN the training_step function already! I can't just simply do this swapping in here, because in the end I would always end up with a single loss, instead I want to have two distinct training steps.
Should I recursively call training_step(...) once more with the inverted batch? Or should I double the dataset and rearrange it (non scalable option)? Or is there a better way to do this?

Thanks everyone for reading up to here, any help is appreciated!