Import nle._pyconverter in nle.__init__ so it can be found on some

systems.
facebookresearch · Apr 14, 2022 · 998e1e8 · 998e1e8
1 parent bff3e2e
commit 998e1e8
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 1 deletion.
diff --git a/nle/__init__.py b/nle/__init__.py
@@ -3,4 +3,7 @@
 # This is needed to register envs
 import nle.env
 
+# This is needed to find module on some systems
+import nle.dataset
+
 from .version import __version__
diff --git a/nle/tests/test_converter.py b/nle/tests/test_converter.py
@@ -72,7 +72,7 @@ def convert_n_times(n):
 
         memory_list = memory_usage((convert_n_times, (100,), {}))
         # After warmup the last few iterations should be constant memory
-        memory_array = np.array(memory_list[3:])
+        memory_array = np.array(memory_list[5:])
         memory_change = (memory_array / memory_array[0]) - 1
         assert max(memory_change) < 0.001  # 0.1 per cent