def _fill_diagonal(matrix: Tensor, value: float) -> Tensor:
        """
        Set diagonal elements with the value.

        Args:
            matrix: tensor of shape (p, p)
            value: value that diagonal should be filled with

        Returns:
            modified matrix with inf on diagonal
        """
        p, _ = matrix.shape
        indices = torch.diag(torch.ones(p)).type(torch.bool)
        matrix[indices] = value
        return matrix

        Set diagonal elements with the value.

        Args:
            matrix: tensor of shape (p, p)
            value: value that diagonal should be filled with

        Returns:
            modified matrix with inf on diagonal
        
catalyst/contrib/data/sampler_inbatch.py