Commit e2a0700b authored by Allen Lavoie's avatar Allen Lavoie Committed by TensorFlower Gardener
Browse files

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
parent 3b7429d7
Loading
Loading
Loading
Loading
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please to comment