diff --git a/enumfields/fields.py b/enumfields/fields.py
index 29a77b2..f3989ae 100644
--- a/enumfields/fields.py
+++ b/enumfields/fields.py
@@ -130,7 +130,8 @@ def refresh_from_db(instance, using=None, fields=None, *args, **kwargs):
         cls.refresh_from_db = refresh_from_db
 
         setattr(cls, name, CastOnAssignDescriptor(self))
-        post_save.connect(update_initial_field, sender=cls, weak=False)
+        self._update_initial_field = update_initial_field
+        post_save.connect(update_initial_field, sender=cls)
 
     def to_python(self, value):
         if isinstance(value, self.enum):
diff --git a/setup.py b/setup.py
index 9f36c21..7d0563c 100644
--- a/setup.py
+++ b/setup.py
@@ -28,7 +28,7 @@ def run_tests(self):
 
 setup(
     name='skip-django-choice-enumfields',
-    version='2.0.0',
+    version='2.0.1',
     author='HZDG, Lubos Matl',
     author_email='matllubos@gmail.com',
     description='Real Python Enums for Django.',
diff --git a/tests/test_migration_state.py b/tests/test_migration_state.py
new file mode 100644
index 0000000..a21da6b
--- /dev/null
+++ b/tests/test_migration_state.py
@@ -0,0 +1,17 @@
+import gc
+import weakref
+
+from django.db.migrations.state import ModelState, ProjectState
+
+from .models import MyModel
+
+
+def test_rendered_migration_state_is_released():
+    project_state = ProjectState()
+    project_state.add_model(ModelState.from_model(MyModel))
+    rendered_model = weakref.ref(project_state.apps.get_model('tests', 'MyModel'))
+
+    del project_state
+    gc.collect()
+
+    assert rendered_model() is None
diff --git a/tox.ini b/tox.ini
index e4d2e21..4317657 100644
--- a/tox.ini
+++ b/tox.ini
@@ -4,6 +4,10 @@ envlist =
 
 [testenv]
 setenv = PYTHONPATH = {toxinidir}
-commands = python setup.py test
+commands = pytest tests
 deps =
+    pytest
+    pytest-django
+    djangorestframework
+    pytz
     django4.2: Django>=4.2,<5.0