Fix variable tracking for Templates assigned to subclassed Model attributes
Previously a Template's variables didn't show up in Model.weights, despite being available in Template.weights, Template.trainable_weights, and so on. Like other variable containers, these are now gathered (and forwarded through lists and dictionaries). Also switches gathering to Model._layers from Model.layers (and likewise for CheckpointableDataStructures), since the un-filtered ._layers includes variables added directly to a container type (e.g. a Template). Updates and losses are still gathered from the post-processed/flattened .layers, which is fine for now since Layers are the only things that have updates or losses. PiperOrigin-RevId: 209212023
Loading
Please sign in to comment