docs: add TrainState docstring with Optimisers API (#1110)

The default TrainState docstring suggests using the version with the Optimisers API, but it doesn't appear in the doc website. Add the docstring manually to the website by adding the method to a @docs block.
LuxDL · Dec 3, 2024 · 277f6ab · 277f6ab
1 parent 7a71c24
commit 277f6ab
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/docs/src/api/Lux/utilities.md b/docs/src/api/Lux/utilities.md
@@ -13,6 +13,7 @@ basic building blocks which can be seamlessly composed to create complex trainin
 
 ```@docs
 Training.TrainState
+Training.TrainState(::AbstractLuxLayer, ::Any, ::Any, ::Optimisers.AbstractRule)
 Training.compute_gradients
 Training.apply_gradients
 Training.apply_gradients!