fix lint F401 warning

microsoft · Dec 19, 2024 · bf780d0 · bf780d0
1 parent c5b87ac
commit bf780d0
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/superbench/benchmarks/model_benchmarks/__init__.py b/superbench/benchmarks/model_benchmarks/__init__.py
@@ -13,9 +13,12 @@
 from superbench.benchmarks.model_benchmarks.megatron_gpt3 import MegatronGPT
 from superbench.benchmarks.model_benchmarks.pytorch_llama import PytorchLlama
 
-__all__ = ['ModelBenchmark', 'PytorchBERT', 'PytorchGPT2', 'PytorchCNN', 'PytorchLSTM', 'MegatronGPT', 'PytorchLlama']
-
 # Check for Python version > 3.7 and conditionally import PytorchMixtral
+PytorchMixtral = None
 if sys.version_info > (3, 7):
     from superbench.benchmarks.model_benchmarks.pytorch_mixtral import PytorchMixtral
+
+__all__ = ['ModelBenchmark', 'PytorchBERT', 'PytorchGPT2', 'PytorchCNN', 'PytorchLSTM', 'MegatronGPT', 'PytorchLlama']
+
+if PytorchMixtral is not None:
     __all__.append('PytorchMixtral')