diff --git a/node_wrangler/__init__.py b/node_wrangler/__init__.py
index afce087eb..303da7b29 100644
--- a/node_wrangler/__init__.py
+++ b/node_wrangler/__init__.py
@@ -44,7 +44,7 @@ def register():
         name="Source Socket!",
         default=0,
         description="An internal property used to store the source socket in a Lazy Connect operation")
-    bpy.types.NodeSocketInterface.NWViewerSocket = BoolProperty(
+    bpy.types.NodeTreeInterfaceSocket.NWViewerSocket = BoolProperty(
         name="NW Socket",
         default=False,
         description="An internal property used to determine if a socket is generated by the addon")
@@ -64,4 +64,4 @@ def unregister():
     del bpy.types.Scene.NWLazySource
     del bpy.types.Scene.NWLazyTarget
     del bpy.types.Scene.NWSourceSocket
-    del bpy.types.NodeSocketInterface.NWViewerSocket
+    del bpy.types.NodeTreeInterfaceSocket.NWViewerSocket