@@ -37,8 +37,8 @@ jobs:
       - name: Install TensorFlow
         shell: bash -el {0}  # Important: activates the conda environment
         if: runner.os != 'macOS'
-        run: |
-          pip install tensorflow tensorpack==0.11 tf_slim==1.1.0 tf-keras wrapt==1.14.1
+        run: |  # wrapt==1.14.1
+          pip install tensorflow tensorpack==0.11 tf_slim==1.1.0 tf-keras
 
       - name: Install TensorFlow on macOS
         shell: bash -el {0}  # Important: activates the conda environment
 
@@ -12,12 +12,12 @@
 # Licensed under GNU Lesser General Public License v3.0
 #
 
+from . import _tf_legacy
+
 # Suppress tensorflow warning messages
 import tensorflow as tf
 tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
 
-from . import _tf_legacy
-
 from deeplabcut.pose_estimation_tensorflow.config import *
 from deeplabcut.pose_estimation_tensorflow.datasets import *
 from deeplabcut.pose_estimation_tensorflow.default_config import *
 
@@ -2,6 +2,7 @@
 import sys
 
 os.environ.setdefault("TF_USE_LEGACY_KERAS", "1")
+os.environ.setdefault("WRAPT_DISABLE_EXTENSIONS", "1")
 
 try:
     import tf_keras.src.legacy_tf_layers as legacy_tf_layers