Multiple separate distinct return statements in a forward pass

Is it possible to have multiple distinct return paths in a forward pass? I’ve seen in other posts whenever there is branching logic, they will return multiple values from the same return, kind of like

def forward():

return value1, value2

Instead, I want to be able to have something like

def forward():

if condition:
return value1
… (potentially more layers)
if condition:
return value2

return value

Can the autograd engine support this? How would you set up your losses to know where you exited from and what layers to update? Apologies if this is documented somewhere but I haven’t been able to find it.

Yes, Autograd won’t have a problem since it tracks which ops and parameters were used in the forward pass.

This is use case specific and you might be able to reuse the same condition from the forward.