Method register_buffer

diff2flow/ddim.py:59–63

(self, name, attr)

        self.schedule = schedule

    def register_buffer(self, name, attr):
        if type(attr) == torch.Tensor:
            if attr.device != self.device:
                attr = attr.to(self.device)
        setattr(self, name, attr)

    def make_schedule(self, ddim_num_steps, device, ddim_discretize="uniform", ddim_eta=0., verbose=False):
        self.device = device