ValueError: Found unexpected Arrays on value of type <class 'list'> in static attribute 'blocks' of Pytree type '<class 'utils.nn.STTransformer'>'. This is an error starting from Flax version 0.12.0.
Consider one of the following options:
1. If the attribute is meant to be static, either remove the Array value or wrap it in a static container.
2. Wrap the value with nnx.data on assignment:
_.blocks = nnx.data(...)
3. Annotate the class attribute with nnx.Data:
class STTransformer(Module):
blocks: nnx.Data[list]
4. If the container is a list or dict, try using nnx.List(...) or nnx.Dict(...) instead.
5. Disable pytree for this class:
class STTransformer(Module, pytree=False):
