Solved by vainaijr in post #2 
              
              
                do you mean something like this? 
import torch.nn as nn, torch

class A(nn.Module):
  def __init__(self):
    super().__init__()
    self.lin = nn.Linear(3, 3)
  def forward(self, x):
    return self.lin(x)

class B(nn.Module):
  def __init__(self):
    super().__init__()
    self.lin = nn.Linear(3,…