diff --git a/src/frontend/src/CustomNodes/hooks/use-handle-node-class.tsx b/src/frontend/src/CustomNodes/hooks/use-handle-node-class.tsx
index fd86f51d1..baf78a781 100644
--- a/src/frontend/src/CustomNodes/hooks/use-handle-node-class.tsx
+++ b/src/frontend/src/CustomNodes/hooks/use-handle-node-class.tsx
@@ -1,6 +1,7 @@
 import useFlowStore from "@/stores/flowStore";
 import { NodeType } from "@/types/flow";
 import { cloneDeep } from "lodash";
+import { useUpdateNodeInternals } from "reactflow";
 
 const useHandleNodeClass = (
   nodeId: string,
@@ -10,6 +11,7 @@ const useHandleNodeClass = (
   ) => void,
 ) => {
   const setNode = setMyNode ?? useFlowStore((state) => state.setNode);
+  const updateNodeInternals = useUpdateNodeInternals();
 
   const handleNodeClass = (newNodeClass, type?: string) => {
     setNode(nodeId, (oldNode) => {
@@ -23,6 +25,8 @@ const useHandleNodeClass = (
         newNode.data.type = type;
       }
 
+      updateNodeInternals(nodeId);
+
       return newNode;
     });
   };