Method _wrap_model

mogen/models/utils/gaussian_diffusion.py:1277–1280

(self, model)

                                       **kwargs)

    def _wrap_model(self, model):
        if isinstance(model, _WrappedModel):
            return model
        return _WrappedModel(model, self.timestep_map, self.original_num_steps)


class _WrappedModel: