fix mix transformer tests

keras-team · Dec 1, 2023 · 77c5b66 · 77c5b66
1 parent 124c79a
commit 77c5b66
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/keras_cv/backend/config.py b/keras_cv/backend/config.py
@@ -23,6 +23,7 @@
 if hasattr(keras, "version") and keras.version().startswith("3."):
     _USE_KERAS_3 = True
 
+
 def detect_if_tensorflow_uses_keras_3():
     # We follow the version of keras that tensorflow is configured to use.
     try:

diff --git a/keras_cv/models/backbones/mix_transformer/mix_transformer_backbone.py b/keras_cv/models/backbones/mix_transformer/mix_transformer_backbone.py
@@ -28,7 +28,6 @@
 from keras_cv import layers as cv_layers
 from keras_cv.api_export import keras_cv_export
 from keras_cv.backend import keras
-from keras_cv.backend import ops
 from keras_cv.models import utils
 from keras_cv.models.backbones.backbone import Backbone
 from keras_cv.models.backbones.mix_transformer.mix_transformer_backbone_presets import (  # noqa: E501
@@ -140,8 +139,8 @@ def __init__(
             # call in `OverlappingPatchingAndEmbedding`
             stride = 4 if i == 0 else 2
             new_height, new_width = (
-                int(ops.shape(x)[1] / stride),
-                int(ops.shape(x)[2] / stride),
+                int(x.shape[1] / stride),
+                int(x.shape[2] / stride),
             )
 
             x = patch_embedding_layers[i](x)