fix onnx exporting

leondgarse · Jun 1, 2023 · 2cbcb97 · 2cbcb97
1 parent 342fa8f
commit 2cbcb97
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/keras_cv_attention_models/__init__.py b/keras_cv_attention_models/__init__.py
@@ -13,6 +13,7 @@
 from keras_cv_attention_models import coat
 from keras_cv_attention_models import coatnet
 from keras_cv_attention_models import convnext
+from keras_cv_attention_models import cotnet
 from keras_cv_attention_models import cmt
 from keras_cv_attention_models import davit
 from keras_cv_attention_models import efficientnet
@@ -23,6 +24,7 @@
 from keras_cv_attention_models import ghostnet
 from keras_cv_attention_models import ghostnet as ghostnetv2  # Will be removed
 from keras_cv_attention_models import gpt2
+from keras_cv_attention_models import halonet
 from keras_cv_attention_models import iformer
 from keras_cv_attention_models import levit
 from keras_cv_attention_models import mlp_family
@@ -68,8 +70,6 @@
 from keras_cv_attention_models import coco
 
 if backend.is_tensorflow_backend:
-    from keras_cv_attention_models import cotnet
-    from keras_cv_attention_models import halonet
     from keras_cv_attention_models import hornet
     from keras_cv_attention_models import nfnets
     from keras_cv_attention_models import volo