diff --git a/mlair/run_modules/model_setup.py b/mlair/run_modules/model_setup.py
index 98263eb732d8067fba0950c7a4882fb3ef020995..2d6d8396b3b22d15ff29083001958b0f4aad244c 100644
--- a/mlair/run_modules/model_setup.py
+++ b/mlair/run_modules/model_setup.py
@@ -202,5 +202,6 @@ class ModelSetup(RunEnvironment):
     @staticmethod
     def _clean_name(orig_name: str):
         mod_name = re.sub(r'^{0}'.format(re.escape("<")), '', orig_name).replace("'", "").split(" ")
-        mod_name = mod_name[1] if any(map(lambda x: x in mod_name[0], ["class", "function", "method"])) else mod_name[0]
-        return mod_name[:-1] if mod_name[-1] == ">" else mod_name
+        mod_name = mod_name[1] if any(map(lambda x: x in mod_name[0], ["class", "function", "method"])) else mod_name
+        mod_name = mod_name[0] if len(mod_name) == 1 else mod_name
+        return mod_name[:-1] if mod_name[-1] == ">" else "".join(mod_name)