Fix import error in Lightning (#5155)

28da51a4 · Yuge Zhang · GitHub · ec5f3e63 · 28da51a4
Unverified Commit 28da51a4 authored Oct 17, 2022 by Yuge Zhang Committed by GitHub Oct 17, 2022
Show whitespace changes
Inline Side-by-side

Showing with 8 additions and 1 deletion

nni/nas/evaluator/pytorch/__init__.py nni/nas/evaluator/pytorch/__init__.py +8 -1

No files found.
--- a/nni/nas/evaluator/pytorch/__init__.py
+++ b/nni/nas/evaluator/pytorch/__init__.py
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT license.

-from .lightning import *
+import warnings
+
+try:
+    from .lightning import *
+except ImportError:
+    warnings.warn("PyTorch-Lightning must be installed to use PyTorch in NAS. "
+                  "If you are not using PyTorch, please `nni.set_default_framework('none')`")
+    raise