Track the rng state.
(self, name, seed)
| 170 | self.states_ = states |
| 171 | |
| 172 | def add(self, name, seed): |
| 173 | """Track the rng state.""" |
| 174 | # Check seed is not already used. |
| 175 | if seed in self.seeds_: |
| 176 | raise Exception("seed {} already exists".format(seed)) |
| 177 | self.seeds_.add(seed) |
| 178 | # Check that state is not already defined. |
| 179 | if name in self.states_: |
| 180 | raise Exception("cuda rng state {} already exists".format(name)) |
| 181 | # Get the current rng state. |
| 182 | orig_rng_state = torch.cuda.get_rng_state() |
| 183 | # Set the new state and store it. |
| 184 | torch.cuda.manual_seed(seed) |
| 185 | self.states_[name] = torch.cuda.get_rng_state() |
| 186 | # Reset rng state to what it was. |
| 187 | _set_cuda_rng_state(orig_rng_state) |
| 188 | |
| 189 | @contextlib.contextmanager |
| 190 | def fork(self, name=_MODEL_PARALLEL_RNG_TRACKER_NAME): |
no test coverage detected