diff --git a/src/backend/langflow/graph/vertex/base.py b/src/backend/langflow/graph/vertex/base.py
index 04dadab85..0ea2e0779 100644
--- a/src/backend/langflow/graph/vertex/base.py
+++ b/src/backend/langflow/graph/vertex/base.py
@@ -48,9 +48,10 @@ class Vertex:
         ]
 
         template_dict = self.data["node"]["template"]
-        self.vertex_type = (
-            self.data["type"] if "Tool" not in self.output else template_dict["_type"]
-        )
+        # self.vertex_type = (
+        #     self.data["type"] if "Tool" not in self.output else template_dict["_type"]
+        # )
+        self.vertex_type = template_dict["_type"]
         if self.base_type is None:
             for base_type, value in ALL_TYPES_DICT.items():
                 if self.vertex_type in value:
diff --git a/src/backend/langflow/interface/loading.py b/src/backend/langflow/interface/loading.py
index d0ec4b845..f8d26ffca 100644
--- a/src/backend/langflow/interface/loading.py
+++ b/src/backend/langflow/interface/loading.py
@@ -32,8 +32,7 @@ def instantiate_class(node_type: str, base_type: str, params: Dict) -> Any:
     params = convert_params_to_sets(params)
     params = convert_kwargs(params)
     if node_type in CUSTOM_AGENTS:
-        custom_agent = CUSTOM_AGENTS.get(node_type)
-        if custom_agent:
+        if custom_agent := CUSTOM_AGENTS.get(node_type):
             return custom_agent.initialize(**params)
 
     class_object = import_by_type(_type=base_type, name=node_type)