Update torchtnt for Python 3.12 (#964)

Summary:
Pull Request resolved: #964

In Python 3.12, dataclass field defaults must be immutable, and implementing a custom `__eq__` on the enum will cause an error like:
```
ValueError: mutable default <enum 'StoppingMechanism'> for field stopping_mechanism is not allowed: use default_factory
```

To work around this, we can inherit from `StrEnum` (newly added into Python 3.12) which supports direct `==` comparisons against strings. For backwards compatibility, we define our own version of it as well.

Reviewed By: diego-urgell

Differential Revision: D68466161

fbshipit-source-id: 1c831cb4449adc953d921ee3559616d5c7b00f86

Author: alanhdu

torchtnt/utils/data/iterators.py

@@ -10,11 +10,9 @@
 from __future__ import annotations
 
 import logging
-
 import random
 from abc import abstractmethod
 from dataclasses import dataclass
-from enum import Enum
 from itertools import cycle
 from typing import (
     Any,
@@ -31,6 +29,16 @@
     Union,
 )
 
+try:
+    # pyre-ignore[21]: Could not find name `StrEnum` in `enum`
+    from enum import StrEnum
+except ImportError:
+    from enum import Enum
+
+    class StrEnum(str, Enum):
+        pass
+
+
 import torch
 import torch.distributed as dist
 
@@ -86,23 +94,14 @@ def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
         pass
 
 
-class StoppingMechanism(Enum):
+class StoppingMechanism(StrEnum):
     ALL_DATASETS_EXHAUSTED = "ALL_DATASETS_EXHAUSTED"
     SMALLEST_DATASET_EXHAUSTED = "SMALLEST_DATASET_EXHAUSTED"
     RESTART_UNTIL_ALL_DATASETS_EXHAUSTED = "RESTART_UNTIL_ALL_DATASETS_EXHAUSTED"
 
     # 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):