From 17ba37d281548e5d60215f741888066717ad5b3e Mon Sep 17 00:00:00 2001 From: Joe Cummings <jrcummings27@gmail.com> Date: Wed, 23 Oct 2024 14:02:20 +0100 Subject: [PATCH] Correct `torchao` check for `TensorCoreTiledLayout` (#1886) --- torchtune/utils/_import_guard.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/torchtune/utils/_import_guard.py b/torchtune/utils/_import_guard.py index 93e7941fbc..06904d9cb6 100644 --- a/torchtune/utils/_import_guard.py +++ b/torchtune/utils/_import_guard.py @@ -20,7 +20,7 @@ _USE_NEW_TENSOR_CORE_TILED_LAYOUT_API = _is_fbcode() or ( not _is_fbcode() and ( - ("dev" not in torchao_version and torchao_version >= "0.6.0") + ("dev" not in torchao_version and torchao_version >= "0.7.0") or ( "dev" in torchao_version and _nightly_version_ge(torchao_version, "2024-10-10")