diff --git a/skops/io/_visualize.py b/skops/io/_visualize.py
index 953345e7..75269bf0 100644
--- a/skops/io/_visualize.py
+++ b/skops/io/_visualize.py
@@ -347,7 +347,8 @@ def visualize(
 
     with zf as zip_file:
         schema = json.loads(zip_file.read("schema.json"))
-        tree = get_tree(schema, load_context=LoadContext(src=zip_file))
+        load_context = LoadContext(src=zip_file, protocol=schema["protocol"])
+        tree = get_tree(schema, load_context=load_context)
 
     nodes = walk_tree(tree)
     # TODO: it would be nice to print html representation if inside a notebook
diff --git a/skops/io/tests/test_audit.py b/skops/io/tests/test_audit.py
index 539c4545..8bd68867 100644
--- a/skops/io/tests/test_audit.py
+++ b/skops/io/tests/test_audit.py
@@ -168,7 +168,7 @@ def test_complex_pipeline_untrusted_set():
 def test_format_object_node():
     estimator = LogisticRegression(random_state=0, solver="liblinear")
     state = get_state(estimator, SaveContext(None))
-    node = ObjectNode(state, LoadContext(None))
+    node = ObjectNode(state, LoadContext(None, -1))
     expected = "sklearn.linear_model._logistic.LogisticRegression"
     assert node.format() == expected
 
@@ -186,5 +186,5 @@ def test_format_object_node():
 )
 def test_format_json_node(inp, expected):
     state = get_state(inp, SaveContext(None))
-    node = JsonNode(state, LoadContext(None))
+    node = JsonNode(state, LoadContext(None, -1))
     assert node.format() == expected
diff --git a/skops/io/tests/test_visualize.py b/skops/io/tests/test_visualize.py
index 62aa51b0..91dab864 100644
--- a/skops/io/tests/test_visualize.py
+++ b/skops/io/tests/test_visualize.py
@@ -106,7 +106,7 @@ def sink(nodes_iter, *args, **kwargs):
         # case this test needs to be changed.
         assert len(nodes_self_unsafe) == 2
         assert nodes_self_unsafe[0].val == "numpy.int64"
-        assert nodes_self_unsafe[1].val == "test_visualize.function"
+        assert nodes_self_unsafe[1].val == "test_visualize.unsafe_function"
 
         # it's not easy to test the number of indirectly unsafe nodes, because
         # it will depend on the nesting structure; we can only be sure that it's