(self, base: nn.Linear)
| 19 | """ |
| 20 | |
| 21 | def __init__(self, base: nn.Linear): |
| 22 | super().__init__() |
| 23 | self.base = base |
| 24 | # self.register_buffer('activation_norms', torch.zeros([base.in_features], device=self.base.weight.device, requires_grad=False)) |
| 25 | self.activation_norms = torch.zeros( |
| 26 | [base.in_features], device=self.base.weight.device, requires_grad=False |
| 27 | ) |
| 28 | self.n_samples = 0 |
| 29 | self.record_activation = True |
| 30 | |
| 31 | def clear_act_buffer(self): |
| 32 | self.activation_norms.fill_(0.0) |