diff --git a/src/backend/langflow/graph/vertex/types.py b/src/backend/langflow/graph/vertex/types.py
index 42a1ee316..868cc6c67 100644
--- a/src/backend/langflow/graph/vertex/types.py
+++ b/src/backend/langflow/graph/vertex/types.py
@@ -131,6 +131,7 @@ class VectorStoreVertex(Vertex):
         super().__init__(data, graph=graph, base_type="vectorstores")
 
         self.params = params or {}
+        self.params.pop("code", None)
 
     # VectorStores may contain databse connections
     # so we need to define the __reduce__ method and the __setstate__ method
@@ -200,9 +201,6 @@ class ChainVertex(Vertex):
             self.params.pop("code", None)
             # Check if the chain requires a PromptVertex
 
-            # Temporarily remove the code from the params
-            self.params.pop("code", None)
-
             for key, value in self.params.items():
                 if isinstance(value, PromptVertex):
                     # Build the PromptVertex, passing the tools if available