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