how do yo compute the final loss (the one where backward is launched from)? I assume both end and out contribute to that loss? If so, this looks OK to me.
This should be an issue for your current use case,  but I want to mention that this probably won’t work correctly with find_unused_parameters=True mode. Because the mlp is used outside of forward, and DDP will find unused parameters using forward output. So in that mode, DDP would treat parameters in mlp as unused parameters although they are actually part of the autograd graph.