from typing import List
import torch


def _collate_slices(
    slices: List[torch.Tensor],
) -> torch.Tensor:
    """
    Convert a list of 2D slices into a padded 3D tensor
    Args:
        slices (list): list of 2D slices of size L[i]*f_dim. Where L[i] is
            length of i-th frame and f_dim is static dimension of features
    Returns:
        3D tensor of size len(slices)*len_max*f_dim where len_max is max of L[i]
    """
    batch_size = len(slices)
    max_len = max(slice.size(0) for slice in slices)

    # for 3 channels: image_dims = [channels x height x width]
    image_dims = slices[0].size()[1:]
    out = slices[0].new_zeros(batch_size, max_len, *image_dims)
    for i, v in enumerate(slices):
        out[i, : v.size(0)] = v
    return out