def forward(
        self, x0: torch.Tensor, x1: torch.Tensor, return_features: bool = False
    ):
        """Symmetrizes the forward pass (see _forward).

        Performs two forward passes, once where x0 is passed through the encoder
        and x1 through the momentum encoder and once the other way around.

        Note that this model currently requires two inputs for the forward pass
        (x0 and x1) which correspond to the two augmentations.
        Furthermore, `the return_features` argument does not work yet.

        Args:
            x0:
                Tensor of shape bsz x channels x W x H.
            x1:
                Tensor of shape bsz x channels x W x H.

        Returns:
            A tuple out0, out1, where out0 and out1 are tuples containing the
            predictions and projections of x0 and x1: out0 = (z0, p0) and
            out1 = (z1, p1).

        Examples:
            >>> # initialize the model and the loss function
            >>> model = BYOL()
            >>> criterion = SymNegCosineSimilarityLoss()
            >>>
            >>> # forward pass for two batches of transformed images x1 and x2
            >>> out0, out1 = model(x0, x1)
            >>> loss = criterion(out0, out1)

        """

        if x0 is None:
            raise ValueError("x0 must not be None!")
        if x1 is None:
            raise ValueError("x1 must not be None!")

        if not all([s0 == s1 for s0, s1 in zip(x0.shape, x1.shape)]):
            raise ValueError(
                f"x0 and x1 must have same shape but got shapes {x0.shape} and {x1.shape}!"
            )

        p0, z1 = self._forward(x0, x1)
        p1, z0 = self._forward(x1, x0)

        return (z0, p0), (z1, p1)
Symmetrizes the forward pass (see _forward).

        Performs two forward passes, once where x0 is passed through the encoder
        and x1 through the momentum encoder and once the other way around.

        Note that this model currently requires two inputs for the forward pass
        (x0 and x1) which correspond to the two augmentations.
        Furthermore, `the return_features` argument does not work yet.

        Args:
            x0:
                Tensor of shape bsz x channels x W x H.
            x1:
                Tensor of shape bsz x channels x W x H.

        Returns:
            A tuple out0, out1, where out0 and out1 are tuples containing the
            predictions and projections of x0 and x1: out0 = (z0, p0) and
            out1 = (z1, p1).

        Examples:
            >>> # initialize the model and the loss function
            >>> model = BYOL()
            >>> criterion = SymNegCosineSimilarityLoss()
            >>>
            >>> # forward pass for two batches of transformed images x1 and x2
            >>> out0, out1 = model(x0, x1)
            >>> loss = criterion(out0, out1)

        
lightly/models/byol.py