Forward training param

code_examples/2_advanced/1_implementing_custom_modules_and_programs_via_subclassing.py

@@ -150,10 +150,16 @@ def _(synalinks):
                 # (don't run the module if no inputs provided)
                 return None
             for i in self.max_iterations:
-                thinking = await self.thinking(inputs)
+                thinking = await self.thinking(
+                    inputs,
+                    training=training,
+                )
                 reward = 0.0
                 if self.critique_program:
-                    critique = await self.critique_program(thinking)
+                    critique = await self.critique_program(
+                        thinking,
+                        training=training,
+                    )
                     reward = critique.get("reward")
                 else:
                     reward = thinking.get("reward")
@@ -165,7 +171,10 @@ def _(synalinks):
                         thinking,
                         name=self.name + f"_inputs_with_thinking_{i}",
                     )
-            return await self.generator(inputs)
+            return await self.generator(
+                inputs,
+                training=training,
+            )
 
         async def compute_output_spec(self, _, training=False):
             return synalinks.SymbolicDataModel(self.schema)