def METHOD_NAME(func):
    def wrapper(self, normalize=False, encoder_out=None, *args, **kwargs):
        if self.ensemble_models is None or len(self.ensemble_models) == 1:
            return func(
                self, normalize=normalize, encoder_out=encoder_out, *args, **kwargs
            )
        def _replace(encoder_out, new_val):
            new_encoder_out = encoder_out.copy()
            new_encoder_out["encoder_out"] = [new_val]
            return new_encoder_out
        action_outs = [
            func(
                model,
                normalize=normalize,
                encoder_out=_replace(
                    encoder_out, encoder_out["encoder_out"][0][:, :, :, i]
                ),
                *args,
                **kwargs
            )
            for i, model in enumerate(self.ensemble_models)
        ]
        if not isinstance(action_outs[0], tuple):  # return multiple values
            action_outs = [[a] for a in action_outs]
        else:
            action_outs = [list(a) for a in action_outs]
        ensembled_outs = []
        for i in range(len(action_outs[0])):
            if i == 0 and normalize:
                ensembled_outs += [
                    torch.logsumexp(
                        torch.stack([a[i] for a in action_outs], -1), dim=-1
                    )
                    - math.log(len(self.ensemble_models))
                ]
            elif action_outs[0][i] is not None:
                ensembled_outs += [torch.stack([a[i] for a in action_outs], -1)]
            else:
                ensembled_outs += [None]
        if len(ensembled_outs) == 1:
            return ensembled_outs[0]
        return tuple(ensembled_outs)
    return wrapper