Method model

tests/schedulers/test_schedulers.py:359–366

(sample, t, *args)

    def dummy_model(self):
        def model(sample, t, *args):
            # if t is a tensor, match the number of dimensions of sample
            if isinstance(t, torch.Tensor):
                num_dims = len(sample.shape)
                # pad t with 1s to match num_dims
                t = t.reshape(-1, *(1,) * (num_dims - 1)).to(sample.device, dtype=sample.dtype)

            return sample * t / (t + 1)

        return model