(self, tensor: Tensor, step: int)
| 115 | self._last_ratio = mask.float().mean().item() |
| 116 | |
| 117 | def apply(self, tensor: Tensor, step: int) -> MaskResult: |
| 118 | self._ensure_module(tensor) |
| 119 | modulus = self._module |
| 120 | assert modulus is not None |
| 121 | metrics: Dict[str, Any] = {"mask_strategy": "learnable", "mask_ratio": self._last_ratio} |
| 122 | if step <= self.config.start_step: |
| 123 | return MaskResult(value=tensor, metrics=metrics) |
| 124 | masked = modulus(tensor, step) |
| 125 | current_ratio = modulus.get_mask_ratio().item() |
| 126 | self._last_ratio = current_ratio |
| 127 | metrics["mask_ratio"] = current_ratio |
| 128 | raw_loss = modulus.get_sparsity_loss() |
| 129 | loss = None if self.config.regularization_weight == 0.0 else raw_loss * self.config.regularization_weight |
| 130 | return MaskResult(value=masked, loss=loss, metrics=metrics) |
| 131 | |
| 132 | def step_optimizer(self, step: int) -> None: |
| 133 | if step <= self.config.start_step: |
nothing calls this directly
no test coverage detected