| 477 | logger.info(f'Added {len(dstore)} keys in {time.time() - start_time:.2f} s') |
| 478 | |
| 479 | class ActivationCapturer(nn.Module): |
| 480 | def __init__(self, layer, capture_input=False): |
| 481 | super().__init__() |
| 482 | self.layer = layer |
| 483 | self.capture_input = capture_input |
| 484 | |
| 485 | self.captured = None |
| 486 | |
| 487 | def forward(self, module, input, output): |
| 488 | if self.capture_input: |
| 489 | self.captured = input[0].detach() |
| 490 | else: |
| 491 | self.captured = output.detach() |
| 492 | |
| 493 | |
| 494 | def get_dstore_path(dstore_dir, model_type, eval_subset, dimension): |
no outgoing calls
no test coverage detected