(self)
| 40 | """ |
| 41 | |
| 42 | def __init__(self): |
| 43 | super().__init__() |
| 44 | self.var_names_to_save: T.Set[str] = set() |
| 45 | self.var_names_to_load: T.Set[str] = set() |
| 46 | self.buffer_names: T.Set[str] = set() |
| 47 | self.parameter_names: T.Set[str] = set() |
| 48 | self.device = torch.device("cpu") |
| 49 | |
| 50 | def register_buffer(self, name: str, tensor: torch.Tensor, persistent=True): |
| 51 | self.buffer_names.add(name) |