Fix pytype problem.

PiperOrigin-RevId: 665776374
google-research · Aug 21, 2024 · b56c1b0 · b56c1b0
1 parent c4bcd66
commit b56c1b0
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/vmoe/projects/contrastive/evaluators.py b/vmoe/projects/contrastive/evaluators.py
@@ -57,6 +57,12 @@ def __init__(
     if bv_discriminative is None:
       raise NotImplementedError(
           'Big Vision must be installed to run the discriminative evaluation.')
+    if (
+        not hasattr(data_sharding.mesh, 'devices')
+        or data_sharding.mesh.devices is None
+    ):
+      raise ValueError(
+          'data_sharding.mesh must be a Mesh, not an AbstractMesh.')
     bv_evaluator = bv_discriminative.Evaluator(
         predict_fn=apply_fn,
         devices=list(data_sharding.mesh.devices.flatten()),
@@ -120,6 +126,12 @@ def __init__(
     if bv_retrieval is None:
       raise NotImplementedError(
           'Big Vision must be installed to run the retrieval evaluation.')
+    if (
+        not hasattr(data_sharding.mesh, 'devices')
+        or data_sharding.mesh.devices is None
+    ):
+      raise ValueError(
+          'data_sharding.mesh must be a Mesh, not an AbstractMesh.')
     bv_evaluator = bv_retrieval.Evaluator(
         predict_fn=apply_fn,
         devices=list(data_sharding.mesh.devices.flatten()),