Your approach should work as seen here:
class NeuraNetworkClass(nn.Module):
    def __init__(self,):
        super().__init__()
        self.fc1 = nn.Linear(1, 1)

    def forward(self, x):
        return self.fc1(x)


class NonNeuralNetworkClass(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.model = model

    def forward(self, a, b):
        return self.model(a) + self.model(b)

model = NeuraNetworkClass()
net = NonNeuralNetworkClass(model)

optim = torch.optim.Adam(lr=1e-4, params=model.parameters())

loss_func = nn.MSELoss()
for epoch in range(10):
    inp1 = torch.randn(1, 1)
    inp2 = torch.randn(1, 1)
    labels = torch.randn(1, 1)
    loss = loss_func(net(a=inp1, b=inp2), labels)
    optim.zero_grad()
    loss.backward()
    for name, param in model.named_parameters():
        print(name, param.grad)
    optim.step()

so a zero gradient could indicate a vanishing gradient issue.