Turn off torch compilation

Author: mattdangerw

keras_nlp/models/bart/bart_seq_2_seq_lm.py

@@ -207,7 +207,6 @@ def __init__(
             loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
             optimizer=keras.optimizers.Adam(2e-5),
             metrics=[keras.metrics.SparseCategoricalAccuracy()],
-            jit_compile=True,
         )
 
     @classproperty

keras_nlp/models/task.py

@@ -79,6 +79,9 @@ def _check_for_loss_mismatch(self, loss):
             )
 
     def compile(self, optimizer="rmsprop", loss=None, **kwargs):
+        # Temporarily disable jit compilation on torch.
+        if config.backend() == "torch":
+            kwargs["jit_compile"] = False
         self._check_for_loss_mismatch(loss)
         super().compile(optimizer=optimizer, loss=loss, **kwargs)
 

keras_nlp/tests/test_case.py

@@ -141,7 +141,9 @@ def call(self, x):
                         return self.layer(x)
 
             model = TestModel(layer)
-            model.compile(optimizer="sgd", loss="mse", jit_compile=True)
+            # Temporarily disable jit compilation on torch backend.
+            jit_compile = config.backend() != "torch"
+            model.compile(optimizer="sgd", loss="mse", jit_compile=jit_compile)
             model.fit(input_data, output_data, verbose=0)
 
         if config.multi_backend():