Method __call__

mogen/models/utils/gaussian_diffusion.py:1290–1295

(self, x, ts, **kwargs)

        self.original_num_steps = original_num_steps

    def __call__(self, x, ts, **kwargs):
        map_tensor = th.tensor(self.timestep_map,
                               device=ts.device,
                               dtype=ts.dtype)
        new_ts = map_tensor[ts]
        return self.model(x, new_ts, **kwargs)