add and use generic torch version comparator

Summary:
# Context
Before we would have an individual function for each version of pytorch we want to check for. Let's remove this in favor of a single universal function

# This Diff
Pass the version as a string instead and replaces old uses

Differential Revision: D56446382

Author: JKSenthil

tests/utils/test_prepare_module.py

@@ -143,9 +143,7 @@ def _test_prepare_module_ddp_throws_with_compile_params_and_static_graph() -> No
         """
 
         tc = unittest.TestCase()
-        with patch(
-            "torchtnt.utils.version.get_torch_version", return_value=Version("2.0.0")
-        ):
+        with patch("torchtnt.utils.version.is_torch_version_geq", return_value=False):
             with tc.assertRaisesRegex(
                 RuntimeError,
                 "Torch version >= 2.1.0 required",

tests/utils/test_version.py

@@ -49,4 +49,4 @@ def test_get_torch_version(self) -> None:
 
     def test_torch_version_comparators(self) -> None:
         with patch.object(torch, "__version__", "2.0.0a0"):
-            self.assertFalse(version.is_torch_version_geq_2_1())
+            self.assertFalse(version.is_torch_version_geq("2.1.0"))

torchtnt/utils/__init__.py

@@ -74,7 +74,7 @@
 from .version import (
     get_python_version,
     get_torch_version,
-    is_torch_version_geq_2_1,
+    is_torch_version_geq,
     is_windows,
 )
 
@@ -144,7 +144,7 @@
     "TLRScheduler",
     "get_python_version",
     "get_torch_version",
-    "is_torch_version_geq_2_1",
+    "is_torch_version_geq",
     "is_windows",
     "get_pet_launch_config",
     "spawn_multi_process",

torchtnt/utils/prepare_module.py

@@ -41,7 +41,7 @@
 )
 
 from torchtnt.utils.rank_zero_log import rank_zero_warn
-from torchtnt.utils.version import is_torch_version_geq_2_1
+from torchtnt.utils.version import is_torch_version_geq
 
 
 @dataclass
@@ -318,7 +318,7 @@ def prepare_module(
             if (
                 torch_compile_params
                 and strategy.static_graph is True
-                and not is_torch_version_geq_2_1()
+                and not is_torch_version_geq("2.1.0")
             ):
                 raise RuntimeError(
                     "Torch version >= 2.1.0 required for Torch compile + DDP with static graph"

torchtnt/utils/version.py

@@ -56,5 +56,5 @@ def get_torch_version() -> Version:
     return pkg_version
 
 
-def is_torch_version_geq_2_1() -> bool:
-    return get_torch_version() >= Version("2.1.0")
+def is_torch_version_geq(version: str) -> bool:
+    return get_torch_version() >= Version(version)