I think the second solution should work! I can’t see a reason why it wouldn’t, at least. Did you try it?
For nested modules, here’s a simplified example:
from torch import nn

class SubModule(nn.Module):
    def __init__(self):
        super(SubModule, self).__init__()
        self.layer1 = ...
    def forward(self, x):
        return self.layer1(x)

class Module(nn.Module):
    def __init__(self, submodule):
        super(Module, self).__init__()
        self.submodule = submodule
        self.layer2 = ...
    def forward(self, x):
        x = self.submodule(x)
        x = self.layer2(x)
        return x