rebase

boomanaiden154 · boomanaiden154 · commit f95f2c8c2eb2 · 2025-04-17T13:05:54.000Z
Created using spr 1.3.4
diff --git a/gematria/model/python/model_base.py b/gematria/model/python/model_base.py
@@ -1306,6 +1306,9 @@ def compute_loss_tensor(self, schedule: FeedDict):
         )
     )
 
+  def _get_trainable_variables(self):
+    return self.trainable_variables
+
   def train_batch(
       self,
       schedule: FeedDict,
@@ -1341,10 +1344,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
`@@ -1306,6 +1306,9 @@ def compute_loss_tensor(self, schedule: FeedDict):`
`1306`	`1306`	`)`
`1307`	`1307`	`)`
`1308`	`1308`
	`1309`	`+ def _get_trainable_variables(self):`
	`1310`	`+ return self.trainable_variables`
	`1311`	`+`
`1309`	`1312`	`def train_batch(`
`1310`	`1313`	`self,`
`1311`	`1314`	`schedule: FeedDict,`
`@@ -1341,10 +1344,11 @@ def train_batch(`
`1341`	`1344`	`]`
`1342`	`1345`	`)`
`1343`	`1346`
	`1347`	`+ trainable_variables = self._get_trainable_variables()`
`1344`	`1348`	`variables = (`
`1345`	`1349`	`[variable.deref() for variable in variables]`
`1346`	`1350`	`if variables`
`1347`		`- else self.trainable_variables`
	`1351`	`+ else trainable_variables`
`1348`	`1352`	`)`
`1349`	`1353`
`1350`	`1354`	`grads = tape.gradient(loss_tensor, variables)`