Thats correct, you would have to implement the logic by yourself. It would look like this:
class MyModule(torch.nn.Module):
    def __init__(self, *args, **kwargs):
        super().__init__()
        self.module_gpu = torch.nn.Conv2d(3, 3, 3, padding=1)
        self.module_cpu = torch.nn.Conv2d(3, 3, 3, padding=1)

    def push_to_devices(self):
        self.module_gpu = self.module_gpu.cuda()
        self.modlue_cpu = self.module_cpu.cpu()

    def forward(self, x):
        x = x.cuda()
        x = self.module_gpu(x)
        x = x.cpu()
        return self._module_cpu(x)


module = MyModule()
in = torch.rand(1, 3, 64, 64)
print(in.device)
module.push_to_devices()
out = module(in)
print(out.device)


Note: You would have to module.push_to_devices() after every change of the networks devices, since this would also change the devices of these layers. If you don’t have to do it like this, I would recommend not to do it, since this usually leads to errors and your code is not device agnostic (i.e. won’t work without a GPU without code changes)