def broadcast_dim_conv2d(x):
    """
    Auto broadcast input so that it can fits into a Conv2d
    """
    if x.dim() == 3:
        x = x[:, (None), :, :]
    else:
        raise ValueError(
            'Only support input with shape = (batch, len) or shape = (len)')
    return x
