Reshape tensor to merge first two dimensions

def merge_first_two_dims(tensor):
    """ Reshape tensor to merge first two dimensions """
    shape = tensor.shape
    batch_size = shape[0] * shape[1]
    new_shape = tuple([batch_size] + list(shape[2:]))
    return tensor.view(new_shape)