rebase

Created using spr 1.3.4

Author: boomanaiden154

gematria/granite/python/gnn_model_base.py

@@ -310,6 +310,12 @@ def initialize(self):
               tf_keras.layers.LayerNormalization(name=globals_layer_norm_name)
           )
 
+  def _get_trainable_variables(self):
+    trainable_variables = list(super()._get_trainable_variables())
+    for layer in self._graph_network:
+      trainable_variables.extend(layer.module.trainable_variables)
+    return trainable_variables
+
   # @Override
   def _forward(self, feed_dict):
     graph_tuple_outputs = self._execute_graph_network(feed_dict)

gematria/model/python/model_base.py

@@ -1312,6 +1312,9 @@ def compute_loss_tensor(self, schedule: FeedDict):
         )
     )
 
+  def _get_trainable_variables(self):
+    return self.trainable_variables
+
   def train_batch(
       self,
       schedule: FeedDict,
@@ -1347,10 +1350,11 @@ def train_batch(
               ]
           )
 
+        trainable_variables = self._get_trainable_variables()
         variables = (
             [variable.deref() for variable in variables]
             if variables
-            else self.trainable_variables
+            else trainable_variables
         )
 
         grads = tape.gradient(loss_tensor, variables)