flax
543686b1 - have TrainState subclass nnx.Module

Commit

2 years ago

have TrainState subclass nnx.Module

References

#3782 - Add `nnx.training`

Author

chiamp

chiamp

Committer

chiamp

chiamp

Parents

Loading