rebase

boomanaiden154 · boomanaiden154 · commit 954c41d131d8 · 2025-04-17T13:06:05.000Z
Created using spr 1.3.4
diff --git a/gematria/model/python/model_base.py b/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)

Original file line number	Diff line number	Diff line change
`@@ -1312,6 +1312,9 @@ def compute_loss_tensor(self, schedule: FeedDict):`
`1312`	`1312`	`)`
`1313`	`1313`	`)`
`1314`	`1314`
	`1315`	`+ def _get_trainable_variables(self):`
	`1316`	`+ return self.trainable_variables`
	`1317`	`+`
`1315`	`1318`	`def train_batch(`
`1316`	`1319`	`self,`
`1317`	`1320`	`schedule: FeedDict,`
`@@ -1347,10 +1350,11 @@ def train_batch(`
`1347`	`1350`	`]`
`1348`	`1351`	`)`
`1349`	`1352`
	`1353`	`+ trainable_variables = self._get_trainable_variables()`
`1350`	`1354`	`variables = (`
`1351`	`1355`	`[variable.deref() for variable in variables]`
`1352`	`1356`	`if variables`
`1353`		`- else self.trainable_variables`
	`1357`	`+ else trainable_variables`
`1354`	`1358`	`)`
`1355`	`1359`
`1356`	`1360`	`grads = tape.gradient(loss_tensor, variables)`