Update torch eval mode

keras/src/backend/jax/trainer.py

@@ -536,7 +536,7 @@ def evaluate(
                 steps_per_execution=self.steps_per_execution,
             )
 
-        self._symbolic_build(iterator=epoch_iterator, training=False)
+        self._symbolic_build(iterator=epoch_iterator)
 
         # Container that configures and calls callbacks.
         if not isinstance(callbacks, callbacks_module.CallbackList):
@@ -765,7 +765,7 @@ def test_on_batch(
         data = (x, y, sample_weight)
         data = _distribute_data(data)
         # Maybe build model
-        self._symbolic_build(data_batch=data, training=False)
+        self._symbolic_build(data_batch=data)
         self._record_training_state_sharding_spec()
         self.make_test_function()
 

keras/src/backend/numpy/trainer.py

@@ -91,7 +91,7 @@ def multi_predict_steps(data):
 
         self.predict_function = predict_step
 
-    def _symbolic_build(self, data_batch, training=True):
+    def _symbolic_build(self, data_batch):
         model_unbuilt = not all(layer.built for layer in self._flatten_layers())
         compile_metrics_unbuilt = (
             self._compile_metrics is not None
@@ -113,7 +113,7 @@ def to_symbolic_input(v):
             ) = data_adapter_utils.unpack_x_y_sample_weight(data_batch)
             # Build all model state with `backend.compute_output_spec`.
             try:
-                y_pred = backend.compute_output_spec(self, x, training=training)
+                y_pred = backend.compute_output_spec(self, x)
             except:
                 raise RuntimeError(
                     "Unable to automatically build the model. "
@@ -246,7 +246,7 @@ def evaluate(
             # Build the model on one batch of data.
             for _, data in epoch_iterator.enumerate_epoch():
                 data_batch = data[0]
-                self._symbolic_build(data_batch, training=False)
+                self._symbolic_build(data_batch)
                 break
 
         # Container that configures and calls callbacks.
@@ -304,7 +304,7 @@ def test_on_batch(
         data = (x, y, sample_weight)
 
         # Maybe build model
-        self._symbolic_build(data, training=False)
+        self._symbolic_build(data)
         self.make_test_function()
 
         logs = self.test_function([data])

keras/src/backend/torch/trainer.py

@@ -344,7 +344,7 @@ def evaluate(
                 steps_per_execution=self.steps_per_execution,
             )
 
-        self._symbolic_build(iterator=epoch_iterator, training=False)
+        self._symbolic_build(iterator=epoch_iterator)
 
         # Container that configures and calls callbacks.
         if not isinstance(callbacks, callbacks_module.CallbackList):
@@ -483,7 +483,7 @@ def test_on_batch(
         data = (x, y, sample_weight)
 
         # Maybe build model
-        self._symbolic_build(data_batch=data, training=False)
+        self._symbolic_build(data_batch=data)
         self.make_test_function()
 
         logs = self.test_function([data])

keras/src/trainers/trainer.py

@@ -1000,7 +1000,7 @@ def _assert_compile_called(self, method_name=None):
                 msg += f"calling `{method_name}()`."
             raise ValueError(msg)
 
-    def _symbolic_build(self, iterator=None, data_batch=None, training=True):
+    def _symbolic_build(self, iterator=None, data_batch=None):
         model_unbuilt = not all(layer.built for layer in self._flatten_layers())
         compile_metrics_unbuilt = (
             self._compile_metrics is not None
@@ -1012,12 +1012,10 @@ def _symbolic_build(self, iterator=None, data_batch=None, training=True):
         optimizer_unbuilt = (
             self.optimizer is not None and not self.optimizer.built
         )
-        need_build = model_unbuilt or compile_metrics_unbuilt
-        if backend.backend() != "torch":
-            # TODO: TorchModuleWrapper will have incorrect behavior using
-            # `_symbolic_build`. Not sure why.
-            need_build = need_build or compile_loss_unbuilt
-        if need_build:
+        if model_unbuilt or compile_metrics_unbuilt or compile_loss_unbuilt:
+            if backend.backend() == "torch":
+                original_training = self.training
+                self.eval()
             # Create symbolic tensors matching an input batch.
 
             def to_symbolic_input(v):
@@ -1040,7 +1038,7 @@ def to_symbolic_input(v):
 
             # Build all model state with `backend.compute_output_spec`.
             try:
-                y_pred = backend.compute_output_spec(self, x, training=training)
+                y_pred = backend.compute_output_spec(self, x)
             except Exception as e:
                 raise RuntimeError(
                     "Unable to automatically build the model. "
@@ -1070,8 +1068,10 @@ def to_symbolic_input(v):
                     y,
                     y_pred,
                     sample_weight=sample_weight,
-                    training=training,
                 )
+            if backend.backend() == "torch":
+                if original_training:
+                    self.train()
         if optimizer_unbuilt:
             # Build optimizer
             self.optimizer.build(self.trainable_variables)