From 8175f280e343ca3561d2f1d8d2a434dfe01911b4 Mon Sep 17 00:00:00 2001
From: ZHANG Zhi <850734033@qq.com>
Date: Tue, 2 Feb 2021 08:32:22 +0800
Subject: [PATCH] Fix 'NoneType' Error on jupyter notebooks (#3337)

---
 nni/retiarii/utils.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/nni/retiarii/utils.py b/nni/retiarii/utils.py
index a48f4dc7da..08791a1a96 100644
--- a/nni/retiarii/utils.py
+++ b/nni/retiarii/utils.py
@@ -2,6 +2,7 @@
 import warnings
 from collections import defaultdict
 from typing import Any
+from pathlib import Path
 
 
 def import_(target: str, allow_none: bool = False) -> Any:
@@ -107,7 +108,18 @@ def blackbox_module(cls):
     Register a module. Use it as a decorator.
     """
     frm = inspect.stack()[1]
+
+    assert (inspect.getmodule(frm[0]) is not None), ('unable to locate the definition of the given black box module, '
+                                                     'please define it explicitly in a .py file.')
     module_name = inspect.getmodule(frm[0]).__name__
+
+    if module_name == '__main__':
+        main_file_path = Path(inspect.getsourcefile(frm[0]))
+        if main_file_path.parents[0] != Path('.'):
+            raise RuntimeError(f'you are using "{main_file_path}" to launch your experiment, '
+                               f'please launch the experiment under the directory where "{main_file_path.name}" is located.')
+        module_name = main_file_path.stem
+
     return _blackbox_cls(cls, module_name, 'args')
 
 
@@ -116,6 +128,8 @@ def register_trainer(cls):
     Register a trainer. Use it as a decorator.
     """
     frm = inspect.stack()[1]
+    assert (inspect.getmodule(frm[0]) is not None), ('unable to locate the definition of the given trainer, '
+                                                     'please define it explicitly in a .py file.')
     module_name = inspect.getmodule(frm[0]).__name__
     return _blackbox_cls(cls, module_name, 'full')