diff --git a/coltra/__init__.py b/coltra/__init__.py
index 3f257a6..dddc34b 100644
--- a/coltra/__init__.py
+++ b/coltra/__init__.py
@@ -6,5 +6,5 @@
 from coltra.buffers import Action, Observation
 from coltra import utils
 
-__version__ = "0.2.0"
+__version__ = "0.2.1"
 VERSION = __version__
diff --git a/coltra/research/__init__.py b/coltra/research/__init__.py
index 9a801cd..e69de29 100644
--- a/coltra/research/__init__.py
+++ b/coltra/research/__init__.py
@@ -1 +0,0 @@
-from .policy_fusion import JointModel
diff --git a/tests/research/test_fusion.py b/tests/research/test_fusion.py
index 417f31a..474513c 100644
--- a/tests/research/test_fusion.py
+++ b/tests/research/test_fusion.py
@@ -12,7 +12,7 @@
 from coltra.models import MLPModel
 import pytest
 
-from coltra.research.policy_fusion import JointModel
+from coltra.research.policy_fusion.policy_fusion import JointModel
 
 
 def assert_models_equal(model1: nn.Module, model2: nn.Module):