Merge pull request #604 from ekellbuch:segmenter

PiperOrigin-RevId: 410598120
google · Nov 17, 2021 · 24b796b · 24b796b
2 parents 6bdb31c + fe95e35
commit 24b796b
Show file tree

Hide file tree

Showing 3 changed files with 459 additions and 0 deletions.
diff --git a/uncertainty_baselines/models/__init__.py b/uncertainty_baselines/models/__init__.py
@@ -81,6 +81,16 @@
 except tf.errors.NotFoundError:
   logging.warning('Skipped ViT models due to NotFoundError.', exc_info=True)
 
+# pylint: disable=g-import-not-at-top
+try:
+  # Try to import Segmenter models.
+  from uncertainty_baselines.models.segmenter import segmenter_transformer
+except ImportError:
+  logging.warning('Skipped Segmenter models due to ImportError.', exc_info=True)
+except tf.errors.NotFoundError:
+  logging.warning('Skipped Segmenter models due to NotFoundError.',
+                  exc_info=True)
+
 try:
   # Try to import models depending on tensorflow_models.official.nlp.
   from uncertainty_baselines.models import bert