diff --git a/src/pymgrid/utils/dry_run.py b/src/pymgrid/utils/dry_run.py
index b6577be7..872e43bf 100644
--- a/src/pymgrid/utils/dry_run.py
+++ b/src/pymgrid/utils/dry_run.py
@@ -1,13 +1,10 @@
 import yaml
 
 from contextlib import contextmanager
-from typing import Union
-
-import pymgrid
 
 
 @contextmanager
-def dry_run(pymgrid_object: Union['pymgrid.Microgrid', 'pymgrid.modules.base.BaseMicrogridModule']):
+def dry_run(pymgrid_object):
     """
     A context manager to test modifications of a pymgrid object without modifying said object.