| 97 | self._hysteresis_tracker = self.hysteresis |
| 98 | |
| 99 | def update(self, found_inf): |
| 100 | |
| 101 | # If we have an inf/nan, growth tracker is set to 0 |
| 102 | # and hysterisis tracker is reduced by 1. |
| 103 | if found_inf: |
| 104 | self._growth_tracker = 0 |
| 105 | self._hysteresis_tracker -= 1 |
| 106 | # Now if we are out of hysteresis count, scale down the loss. |
| 107 | if self._hysteresis_tracker <= 0: |
| 108 | self._scale = torch.max( |
| 109 | self._scale * self.backoff_factor, self.min_scale |
| 110 | ) |
| 111 | else: |
| 112 | # If there is no nan/inf, increment the growth tracker. |
| 113 | self._growth_tracker += 1 |
| 114 | # If we have had enough consequitive intervals with no nan/inf: |
| 115 | if self._growth_tracker == self.growth_interval: |
| 116 | # Reset the tracker and hysteresis trackers, |
| 117 | self._growth_tracker = 0 |
| 118 | self._hysteresis_tracker = self.hysteresis |
| 119 | # and scale up the loss scale. |
| 120 | self._scale = self._scale * self.growth_factor |
| 121 | |
| 122 | def state_dict(self): |
| 123 | state_dict = {} |