def METHOD_NAME(self, args):
    # Since hybrid_forward does not support dictionary, we have to separate
    # the raw outputs of the network based on the indices and map them to
    # the learnable parameters
    num_distr_args = len(self.base_distr_args_dim)
    distr_args = args[0:num_distr_args]
    num_transforms_args = [
        len(transform_dim_args)
        for transform_dim_args in self.transforms_args_dim
    ]
    # starting indices of arguments for each transformation
    num_args_cumsum = np.cumsum([num_distr_args] + num_transforms_args)
    # get the arguments for each of the transformations
    transforms_args = list(
        map(
            lambda ixs: args[ixs[0] : ixs[1]],
            zip(num_args_cumsum, num_args_cumsum[1:]),
        )
    )
    return distr_args, transforms_args