I modified it to use ‘up’ or ‘down’ direction
from IPython.core.debugger import set_trace

class Top(nn.Module):
  def __init__(self, direction):
    super().__init__()
    self.encoder = nn.Sequential(nn.Conv2d(1, 1, 1), nn.Dropout(0.5))
    self.lin = nn.Linear(10, 10)
    self.direction = direction
    if self.direction == 'down':
      self.childone = Bottom()
      self.childtwo = Bottom()
  def forward(self, x):
    if self.direction == 'down':
      return self.childone(self.encoder(x)), self.childtwo(self.encoder(x))
    elif self.direction == 'up':
      return self.lin(x)

class Bottom(nn.Module):
  def __init__(self):
    super().__init__()
    self.lin_one = nn.Linear(784, 100)
    self.lin_two = nn.Linear(100, 10)
    self.parent = Top('up')
  def forward(self, x):
    return self.parent(self.lin_two(self.lin_one(x.view(x.size(0), -1))))

but this also has problem, as it gives output to Top class twice, so finally I get tuple instead of one output.