(self, base: nn.Linear)
| 21 | """ |
| 22 | |
| 23 | def __init__(self, base: nn.Linear): |
| 24 | super().__init__() |
| 25 | self.base = base |
| 26 | self.activation_norms = [] # offload to CPU |
| 27 | self.record_activation = True |
| 28 | |
| 29 | def clear_act_buffer(self): |
| 30 | self.activation_norms = [] |