Function test_parallel_transformer_block

tests/tests.py:164–167

(dim, dim_head, heads, ff_mult, random_tensor)

    ],
)
def test_parallel_transformer_block(dim, dim_head, heads, ff_mult, random_tensor):
    block = ParallelTransformerBlock(dim, dim_head, heads, ff_mult)
    output = block(random_tensor)
    assert output.shape == random_tensor.shape


@pytest.mark.parametrize(