def time_distributed(inputs, op, activation=None, **kwargs):
    """
    apply op on both the batch (B) and time (T) dimension
    args:
        inputs: [B, T, ...]
        op: a layer op that accepts [B, ...]
        reshape: if reshape to [B, T, newshape]
    return:
        outputs: [B, T, ...]
    """

    batch_size = inputs.size(0)
    seq_len = inputs.size(1)
    p_dim = inputs.size()[2:]
    outputs = op(inputs.view(-1, *p_dim), **kwargs)
    out_dim = outputs.size()[1:]
    if activation is not None:
        outputs = activation(outputs)
    outputs = outputs.view(batch_size, seq_len, *out_dim)
    return outputs