diff --git a/cellfinder/__init__.py b/cellfinder/__init__.py
index 84b2b242..d830ec72 100644
--- a/cellfinder/__init__.py
+++ b/cellfinder/__init__.py
@@ -27,7 +27,7 @@
     import os
 
     # check if environment variable exists?
-    os.environ["KERAS_BACKEND"] = "jax"  # "torch" "jax", "tensorflow"
+    os.environ["KERAS_BACKEND"] = "torch"  # "torch" "jax", "tensorflow"
 
 except PackageNotFoundError as e:
     raise PackageNotFoundError("error setting up Keras backend") from e
diff --git a/cellfinder/core/train/train_yml.py b/cellfinder/core/train/train_yml.py
index dd83ebc5..05375bcd 100644
--- a/cellfinder/core/train/train_yml.py
+++ b/cellfinder/core/train/train_yml.py
@@ -451,6 +451,9 @@ def run(
         epochs=epochs,
         callbacks=callbacks,
     )
+    # model.built = True,
+    # model.compiled=True, and
+    # compile_metrics_unbuilt ---> I get error
 
     if save_weights:
         logger.info("Saving model weights")