If you want to keep this output during reconstruction-only training, and prevent autograd from storing its intermediate tensors, you can wrap self.layer3(out) execution with torch.set_grad_enabled(False) context manager. In order to swap between your training regimes, you can use the following code example:
class MyNet(nn.Module):
  def __init__(self):
    super().__init__()
    self.seg_train = True
    self.layer1 = nn.Conv2d(3, 32, 3)
    self.layer2 = nn.Conv2d(32, 64, 3)
    self.layer3 = nn.Conv2d(64, 2, 3)
    self.layer4 = nn.Conv2d(64, 1, 3)

  def forward(self, x):
    out = self.layer1(x)
    out = self.layer2(out)

    # in grad is not enabled - we should not enable it
    grad_enabled = torch.is_grad_enabled()
    with torch.set_grad_enabled(self.seg_train and grad_enabled):
        out1 = self.layer3(out)

    out2 = self.layer4(out)

    return out1, out2

Then you will be able to swap training regimes by chaning model.seg_train variable.