diff --git a/test/test_treetensornetworks/test_solvers/Project.toml b/test/test_treetensornetworks/test_solvers/Project.toml
index 91af4fe2..c377be4d 100644
--- a/test/test_treetensornetworks/test_solvers/Project.toml
+++ b/test/test_treetensornetworks/test_solvers/Project.toml
@@ -1,4 +1,5 @@
 [deps]
+DataGraphs = "b5a273c3-7e6c-41f6-98bd-8d7f1525a36a"
 Dictionaries = "85a47980-9c8c-11e8-2b9f-f7ca1fa99fb4"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
 ITensorNetworks = "2919e153-833c-4bdc-8836-1ea460a35fc7"
diff --git a/test/test_treetensornetworks/test_solvers/test_dmrg.jl b/test/test_treetensornetworks/test_solvers/test_dmrg.jl
index 076c533e..91f573ca 100644
--- a/test/test_treetensornetworks/test_solvers/test_dmrg.jl
+++ b/test/test_treetensornetworks/test_solvers/test_dmrg.jl
@@ -1,10 +1,11 @@
 @eval module $(gensym())
+using DataGraphs: vertex_data
 using Dictionaries: Dictionary
 using Graphs: nv, vertices
 using ITensors: ITensors
 using ITensors.ITensorMPS: MPO, MPS
 using ITensorNetworks:
-  ITensorNetworks, OpSum, TTN, apply, dmrg, inner, mpo, random_mps, siteinds
+  ITensorNetworks, OpSum, TTN, apply, dmrg, inner, mpo, random_mps, relabel_sites, siteinds
 using KrylovKit: eigsolve
 using NamedGraphs: named_comb_tree
 using Observers: observer