(self, name, attr)
| 16 | self.schedule = schedule |
| 17 | |
| 18 | def register_buffer(self, name, attr): |
| 19 | if type(attr) == torch.Tensor: |
| 20 | if attr.device != torch.device("cuda"): |
| 21 | attr = attr.to(torch.device("cuda")) |
| 22 | setattr(self, name, attr) |
| 23 | |
| 24 | def make_schedule(self, ddim_num_steps, ddim_discretize="uniform", ddim_eta=0., verbose=True): |
| 25 | if ddim_eta != 0: |