How do I design two networks which part of their layers are same

1

I want to design two networks and update them separately, but some of their layers are the same.

I have two plans. Can you tell me which one is better or give another better idea?

  1. Use two nn.module. Change the weights of the same layers of A network after B network updates.
  2. Use one nn.module. In this situation, A and B can use the same layers.