It seems that x_stageX is undefined in the forward, so the return statement should yield an error.
Based on the description of your use case you would need to define these intermediate tensors either in the forward of your main model or in the submodels. If the latter case is not easily done, you could use e.g. forward hooks as described here.