From 5592267db31fee939ec29bfc5ef282917a80ac38 Mon Sep 17 00:00:00 2001 From: Nihal Harish Date: Tue, 29 Jun 2021 10:55:04 -0700 Subject: [PATCH] cache output of is_framework_version_supported (#508) --- smdebug/core/utils.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/smdebug/core/utils.py b/smdebug/core/utils.py index 6608544a2..54c1c2948 100644 --- a/smdebug/core/utils.py +++ b/smdebug/core/utils.py @@ -7,6 +7,7 @@ import socket import urllib.parse from enum import Enum +from functools import lru_cache from pathlib import Path from typing import Dict, List @@ -638,6 +639,8 @@ def check_smmodelparallel_training(): return _is_using_smmodelparallel +# we need to compute the output of this fn only once since the framework type will remain constant during execution +@lru_cache(maxsize=None) def is_framework_version_supported(framework_type): if framework_type == FRAMEWORK.PYTORCH: from smdebug.pytorch.utils import is_current_version_supported