From 05cb03eb227172f8133f564007da647b4d2b108b Mon Sep 17 00:00:00 2001 From: Benjamin Date: Wed, 26 Jan 2022 17:50:06 -0500 Subject: [PATCH] update install check of transformers to NM_INTEGRATED --- src/sparseml/transformers/__init__.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/sparseml/transformers/__init__.py b/src/sparseml/transformers/__init__.py index 87dc1dfc01f..7d69924ca71 100644 --- a/src/sparseml/transformers/__init__.py +++ b/src/sparseml/transformers/__init__.py @@ -25,7 +25,7 @@ import transformers as _transformers # triggers error if neuralmagic/transformers is not installed - _transformers.models.bert.modeling_bert.QATMatMul + assert _transformers.NM_INTEGRATED _transformers_import_error = None except Exception as _transformers_import_err: _transformers_import_error = _transformers_import_err @@ -101,11 +101,10 @@ def _check_transformers_install(): try: import transformers as _transformers - _transformers.models.bert.modeling_bert.QATMatMul + assert _transformers.NM_INTEGRATED except Exception: _LOGGER.warning( - "transformers.models.bert.modeling_bert.QATMatMul not availalbe. the" - "neuralmagic fork of transformers may not be installed. it can be " + "the neuralmagic fork of transformers may not be installed. it can be " "installed via " f"`pip install {_NM_TRANSFORMERS_NIGHTLY}`" )