fixed imports to also work in package

tueboesen · Mar 16, 2023 · be486f6 · be486f6
1 parent 0766f80
commit be486f6
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 2 deletions.
diff --git a/multibodypendulum/__init__.py b/multibodypendulum/__init__.py
@@ -1 +1,8 @@
-from npendul import MultiBodyPendulum
+from .npendul import MultiBodyPendulum
+
+__version__ = "0.2"
+
+__all__ = [
+    "MultiBodyPendulum",
+    "__version__",
+]
diff --git a/tests/__init__.py b/tests/__init__.py
diff --git a/tests/test_npendulum.py b/tests/test_npendulum.py
@@ -4,7 +4,7 @@
 import pytest
 import torch
 
-from multibodypendulum.npendul import MultiBodyPendulum
+from multibodypendulum import MultiBodyPendulum
 
 
 def test_stability():