(self, name: str, tensor: torch.Tensor, persistent=True)
| 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) |
| 52 | setattr(self, name, tensor) |
| 53 | |
| 54 | if persistent: |
| 55 | self.register_var_to_load(var_name=name) |
| 56 | else: |
| 57 | self.register_var_to_save(var_name=name) |
| 58 | |
| 59 | def register_parameter(self, name: str, param: torch.Tensor): |
| 60 | self.parameter_names.add(name) |
no test coverage detected