def _off_diagonal(x: Tensor) -> Tensor:
    """Returns a flattened view of the off-diagonal elements of a square matrix."""

    # Ensure the input is a square matrix
    n, m = x.shape
    assert n == m

    # Flatten the matrix and extract off-diagonal elements
    return x.flatten()[:-1].view(n - 1, n + 1)[:, 1:].flatten()
Returns a flattened view of the off-diagonal elements of a square matrix.
lightly/loss/barlow_twins_loss.py