Fix incorrect date check for ao nightlies (#1176)

torchtune/modules/low_precision/_utils.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+from datetime import datetime
 from importlib.metadata import PackageNotFoundError, version
 from typing import Optional, Tuple
 
@@ -16,6 +17,18 @@ def _is_fbcode():
     return not hasattr(torch.version, "git_version")
 
 
+def _nightly_version_ge(ao_version_str: str, date: str) -> bool:
+    """
+    Compare a torchao nightly version to a date of the form
+    %Y-%m-%d.
+
+    Returns True if the nightly version is greater than or equal to
+        the date, False otherwise
+    """
+    ao_datetime = datetime.strptime(ao_version_str.split("+")[0], "%Y.%m.%d")
+    return ao_datetime >= datetime.strptime(date, "%Y-%m-%d")
+
+
 def _get_torchao_version() -> Tuple[Optional[str], Optional[bool]]:
     """
     Get torchao version. Returns a tuple of two elements, the first element

torchtune/utils/quantization.py

@@ -13,10 +13,13 @@
     Quantizer,
 )
 
-from torchtune.modules.low_precision._utils import _get_torchao_version
+from torchtune.modules.low_precision._utils import (
+    _get_torchao_version,
+    _nightly_version_ge,
+)
 
 ao_version, is_nightly = _get_torchao_version()
-if is_nightly and (ao_version >= "2024.7.3"):
+if is_nightly and _nightly_version_ge(ao_version, "2024-07-03"):
     from torchao.quantization.quant_api import quantize_ as quantize
 else:
     from torchao.quantization.quant_api import quantize