diff --git a/client/ayon_core/hosts/tvpaint/plugins/load/load_reference_image.py b/client/ayon_core/hosts/tvpaint/plugins/load/load_reference_image.py
index 0a12e93f44..ea670996a4 100644
--- a/client/ayon_core/hosts/tvpaint/plugins/load/load_reference_image.py
+++ b/client/ayon_core/hosts/tvpaint/plugins/load/load_reference_image.py
@@ -210,15 +210,17 @@ class LoadImage(plugin.Loader):
     def switch(self, container, representation):
         self.update(container, representation)
 
-    def update(self, container, representation):
+    def update(self, container, context):
         """Replace container with different version.
 
         New layers are loaded as first step. Then is tried to change data in
         new layers with data from old layers. When that is done old layers are
         removed.
         """
+
+        repre_doc = context["representation"]
         # Create new containers first
-        context = get_representation_context(representation)
+        context = get_representation_context(repre_doc)
 
         # Get layer ids from previous container
         old_layer_names = self.get_members_from_container(container)