def test_plv_time_shape():
    """
    Test that plv_time returns the correct shape.
    Input shape: (batch, channels, time)
    Output shape: (batch, channels, channels)
    """
    batch, channels, time = 2, 4, 500
    input_tensor = torch.randn(batch, channels, time)
    plv_matrix = plv_time(input_tensor, forward_fourier=True)
    expected_shape = (batch, channels, channels)
    assert plv_matrix.shape == expected_shape, f"Expected shape {expected_shape}, got {plv_matrix.shape}"

    Test that plv_time returns the correct shape.
    Input shape: (batch, channels, time)
    Output shape: (batch, channels, channels)
    
test/unit_tests/models/test_functional.py