# Copyright (c) OpenMMLab. All rights reserved.
import torch

from mmaction.models import UniFormer
from mmaction.testing import generate_backbone_demo_inputs


def test_uniformer_backbone():
    """Test uniformer backbone."""
    input_shape = (1, 3, 16, 64, 64)
    imgs = generate_backbone_demo_inputs(input_shape)

    model = UniFormer(
        depth=[3, 4, 8, 3],
        embed_dim=[64, 128, 320, 512],
        head_dim=64,
        drop_path_rate=0.1)
    model.init_weights()

    model.eval()
    assert model(imgs).shape == torch.Size([1, 512, 8, 2, 2])