The following seems to work:
    def forward(self, input: Tensor) -> Tensor:
        out = torch.eye(2)
        out[slice(0, 1), slice(0, 1)] = self.first(input)
        out[slice(1, 2), slice(1, 2)] = self.second(input)
        return out