| 206 | return {'loss': loss} |
| 207 | |
| 208 | def step(self): |
| 209 | # Run the learning step. |
| 210 | results = self._step() |
| 211 | |
| 212 | # Compute elapsed time. |
| 213 | timestamp = time.time() |
| 214 | elapsed_time = timestamp - self._timestamp if self._timestamp else 0 |
| 215 | self._timestamp = timestamp |
| 216 | |
| 217 | # Update our counts and record it. |
| 218 | counts = self._counter.increment(steps=1, walltime=elapsed_time) |
| 219 | results.update(counts) |
| 220 | self._logger.write(results) |
| 221 | |
| 222 | def get_variables(self, names: List[str]) -> List[Variables]: |
| 223 | return [tf2_utils.to_numpy(self._variables)] |