diff --git a/edml/core/start_device.py b/edml/core/start_device.py
index 7894e7b8b352e8a05c34f8d07df5661a69b614cf..49df053b5073a334d185ce21bc30c2188a545716 100644
--- a/edml/core/start_device.py
+++ b/edml/core/start_device.py
@@ -24,10 +24,10 @@ from edml.models.provider.base import ModelProvider
 def _start_device_server(
     device_id,
     address,
-    client: DeviceClient,
-    server: DeviceServer,
-    logger: SimpleLogger,
-    battery: Battery,
+    client: "DeviceClient",
+    server: "DeviceServer",
+    logger: "SimpleLogger",
+    battery: "Battery",
     devices_info,
     max_threads=10,
     max_message_length=-1,
diff --git a/edml/helpers/config_helpers.py b/edml/helpers/config_helpers.py
index 9902777106236e52128e2490817c1c7305bbf5e6..7f4ec87e06ee3542cb28ce3e857880795164188d 100644
--- a/edml/helpers/config_helpers.py
+++ b/edml/helpers/config_helpers.py
@@ -1,12 +1,14 @@
 from copy import deepcopy
 from inspect import signature
+from typing import TYPE_CHECKING
 
 import torch
 from hydra.utils import get_class, instantiate
 from omegaconf import OmegaConf, DictConfig, ListConfig
 from omegaconf.errors import ConfigAttributeError
 
-from edml.controllers.base_controller import BaseController
+if TYPE_CHECKING:
+    from edml.controllers.base_controller import BaseController
 
 
 def get_device_address_by_id(device_id: str, cfg: DictConfig) -> str:
@@ -184,7 +186,7 @@ def drop_irrelevant_keys_recursively(cfg: DictConfig) -> DictConfig:
     return cfg
 
 
-def instantiate_controller(cfg: DictConfig) -> BaseController:
+def instantiate_controller(cfg: DictConfig):  # -> BaseController:
     """
     Instantiates a controller based on the configuration. This method filters out extra parameters defined through hydra
     but not required by the controller's init method. This allows for hydra's multirun feature to work even if