Method
__init__
(self, model, timestep_map, rescale_timesteps,
original_num_steps)
Source from the content-addressed store, hash-verified
| 115 | class _WrappedModel: |
| 116 | |
| 117 | def __init__(self, model, timestep_map, rescale_timesteps, |
| 118 | original_num_steps): |
| 119 | self.model = model |
| 120 | self.timestep_map = timestep_map |
| 121 | self.rescale_timesteps = rescale_timesteps |
| 122 | self.original_num_steps = original_num_steps |
| 123 | |
| 124 | def __call__(self, x, ts, **kwargs): |
| 125 | map_tensor = th.tensor(self.timestep_map, |
Tested by
no test coverage detected