Changed to new model loading in controller

edml/controllers/base_controller.py

@@ -6,7 +6,7 @@ import torch
 
 from edml.controllers.early_stopping import create_early_stopping_callback
 from edml.core.device import DeviceRequestDispatcher
-from edml.helpers.load_model import get_models
+from edml.core.start_device import _get_models
 from edml.helpers.logging import SimpleLogger, create_logger
 from edml.helpers.metrics import ModelMetricResultContainer
 from edml.helpers.types import DeviceBatteryStatus
@@ -41,7 +41,7 @@ class BaseController(abc.ABC):
             )
 
             # if no weights are loaded, initialize the models randomly and set them on all devices
-            client_model, server_model = get_models(self.cfg)
+            client_model, server_model = _get_models(self.cfg)
             self._set_weights_on_all_devices(
                 client_model.state_dict(), on_client=True, wait_for_ready=True
             )