| """Shared base class. | |
| ``torch.optim.Optimizer`` is used strictly as a base class, as the assignment | |
| allows: it supplies parameter-group bookkeeping and the ``state`` dict, and | |
| every update rule below is written out by hand. | |
| """ | |
| import torch | |
| class BaseOptimizer(torch.optim.Optimizer): | |
| def _closure(closure): | |
| if closure is None: | |
| return None | |
| with torch.enable_grad(): | |
| return closure() | |
| def memory_bytes(self) -> int: | |
| """Bytes held in optimizer state -- the number that matters for the | |
| memory-efficient category.""" | |
| total = 0 | |
| for st in self.state.values(): | |
| for v in st.values(): | |
| if torch.is_tensor(v): | |
| total += v.numel() * v.element_size() | |
| return total | |