import torch


def collate_fn(samples):
    """
    collate_fn for SequentialMNIST.
    
    Args:
        samples: a list of samples. Each item is a (imgs, nums) pair. Where
        
            - imgs: shape (T, 1, C, H, W)
            - nums: shape (T, 1)
            
            And len(samples) is the batch size.

    Returns:
        A tuple (imgs, nums). Where
        
        - imgs: shape (T, B, C, H, W)
        - nums: shape (T, B)
    """
    imgs, nums = zip(*samples)
    # (T, B, C, H, W)
    imgs = torch.cat(imgs, dim=1)
    # (T, B)
    nums = torch.cat(nums, dim=1)
    
    return imgs, nums