| 73 | |
| 74 | |
| 75 | class set_mask: |
| 76 | def __init__(self, model, mask): |
| 77 | self.model = model |
| 78 | self.mask = mask |
| 79 | |
| 80 | def __enter__(self): |
| 81 | for name, module in self.model.named_modules(): |
| 82 | if isinstance(module, ActLinear): |
| 83 | module.mask = self.mask |
| 84 | |
| 85 | def __exit__(self, exc_type, exc_val, exc_tb): |
| 86 | for name, module in self.model.named_modules(): |
| 87 | if isinstance(module, ActLinear): |
| 88 | module.mask = None |
| 89 | |
| 90 | |
| 91 | def make_Act(model, verbose=False): |