It does work with submodules.
The problem is you probably have some hardcoded device in the submodules.
This is calling .cuda() instead of coding it as .to(device) being device based on forward’s input.