Contains training state for the learner.
| 48 | |
| 49 | @flax.struct.dataclass |
| 50 | class TrainingState: |
| 51 | """Contains training state for the learner.""" |
| 52 | optimizer_state: optax.OptState |
| 53 | normalizer_params: running_statistics.RunningStatisticsState |
| 54 | cvae_params: Params |
| 55 | |
| 56 | |
| 57 | def _unpmap(v): |
no outgoing calls
no test coverage detected