r"""Update the scale factor according to whether encountered overflow grad. If ``new_scale`` is provided, internal update mechanism will be ignored.
(self, new_scale: float = None)
| 152 | return rst |
| 153 | |
| 154 | def update(self, new_scale: float = None): |
| 155 | r"""Update the scale factor according to whether encountered overflow grad. |
| 156 | If ``new_scale`` is provided, internal update mechanism will be ignored. |
| 157 | """ |
| 158 | if self.growth_interval == 0: |
| 159 | return |
| 160 | |
| 161 | if new_scale is not None: |
| 162 | self.scale_factor = float(new_scale) |
| 163 | else: |
| 164 | if self._found_non_finite: |
| 165 | self.scale_factor *= self.backoff_factor |
| 166 | self._growth_tracker = 0 |
| 167 | else: |
| 168 | self._growth_tracker += 1 |
| 169 | if self._growth_tracker >= self.growth_interval: |
| 170 | self.scale_factor *= self.growth_factor |
| 171 | self._growth_tracker = 0 |
| 172 | self._found_non_finite = False |
| 173 | |
| 174 | def state_dict(self): |
| 175 | return { |