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