def flatten_tensor(inputs):
    """
    Convert a 4D tensor of shape (B, C, H, W) to a 2D tensor of shape 
    (B*H*W, C) and return (B, H, W, C) shape
    """
    inputs = inputs.permute(0, 2, 3, 1).contiguous()
    bhwc = inputs.shape
    return inputs.view(-1, bhwc[-1]), bhwc