Method __init__

mogen/models/utils/gaussian_diffusion.py:1285–1288

(self, model, timestep_map, original_num_steps)

class _WrappedModel:

    def __init__(self, model, timestep_map, original_num_steps):
        self.model = model
        self.timestep_map = timestep_map
        self.original_num_steps = original_num_steps

    def __call__(self, x, ts, **kwargs):
        map_tensor = th.tensor(self.timestep_map,