diff --git a/src/backend/base/langflow/interface/custom/custom_component/component.py b/src/backend/base/langflow/interface/custom/custom_component/component.py
index ce40b0f74..8a9ba1a22 100644
--- a/src/backend/base/langflow/interface/custom/custom_component/component.py
+++ b/src/backend/base/langflow/interface/custom/custom_component/component.py
@@ -34,6 +34,10 @@ class Component:
             if key == "user_id":
                 setattr(self, "_user_id", value)
             else:
+                if key == "code" and "from langflow import CustomComponent" in value:
+                    value = value.replace(
+                        "from langflow import CustomComponent", "from langflow.custom import CustomComponent"
+                    )
                 setattr(self, key, value)
 
     def __setattr__(self, key, value):