Enable comparison of iterator sampler strategy with str value (#954)

Summary: Pull Request resolved: #954

Reviewed By: mpahsu, JKSenthil

Differential Revision: D67550088

fbshipit-source-id: 207b25fde316f9884abcb873fd3c27696efdede5

tests/utils/data/test_iterators.py

@@ -0,0 +1,30 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-strict
+
+import unittest
+
+from torchtnt.utils.data.iterators import StoppingMechanism
+
+
+class TestIterators(unittest.TestCase):
+
+    def test_stopping_mechanism_comparison(self) -> None:
+        self.assertTrue(
+            StoppingMechanism.ALL_DATASETS_EXHAUSTED == "ALL_DATASETS_EXHAUSTED"
+        )
+        self.assertTrue(
+            StoppingMechanism.ALL_DATASETS_EXHAUSTED
+            == StoppingMechanism.ALL_DATASETS_EXHAUSTED
+        )
+        self.assertFalse(
+            StoppingMechanism.ALL_DATASETS_EXHAUSTED == "SMALLEST_DATASET_EXHAUSTED"
+        )
+        self.assertFalse(
+            StoppingMechanism.ALL_DATASETS_EXHAUSTED
+            == StoppingMechanism.SMALLEST_DATASET_EXHAUSTED
+        )

torchtnt/utils/data/iterators.py

@@ -94,6 +94,15 @@ class StoppingMechanism(Enum):
     # used with RandomizedBatchSampler
     WRAP_AROUND_UNTIL_KILLED = "WRAP_AROUND_UNTIL_KILLED"
 
+    def __eq__(self, other: Union[str, StoppingMechanism]) -> bool:
+        """
+        Enable comparison betwen string and instances of StoppingMechanism
+        """
+        if isinstance(other, str):
+            return self.value == other
+
+        return super().__eq__(other)
+
 
 @dataclass
 class RoundRobin(DataIterationStrategy):