diff --git a/Project.toml b/Project.toml
index 1d715cd..bf3858a 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "ITensorTDVP"
 uuid = "25707e16-a4db-4a07-99d9-4d67b7af0342"
 authors = ["Matthew Fishman <mfishman@flatironinstitute.org> and contributors"]
-version = "0.4.3"
+version = "0.4.4"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
@@ -20,7 +20,7 @@ ITensorTDVPObserversExt = "Observers"
 
 [compat]
 Compat = "4"
-ITensors = "0.3.58, 0.4, 0.5, 0.6"
+ITensors = "0.6.10"
 KrylovKit = "0.6, 0.7"
 Observers = "0.2"
 PackageExtensionCompat = "1"
diff --git a/src/expand.jl b/src/expand.jl
index 0d3c62a..aabfa51 100644
--- a/src/expand.jl
+++ b/src/expand.jl
@@ -90,8 +90,7 @@ function expand(
     projectorⱼ = idⱼ - prime(basisⱼ, rinds) * dag(basisⱼ)
     # Sum reference density matrices
     ρⱼ = sum(reference -> prime(reference[j], rinds) * dag(reference[j]), references)
-    # TODO: Fix bug that `tr` isn't preserving the element type.
-    ρⱼ /= scalartype(state)(tr(ρⱼ))
+    ρⱼ /= tr(ρⱼ)
     # Apply projectorⱼ
     ρⱼ_projected = apply(apply(projectorⱼ, ρⱼ), projectorⱼ)
     expanded_basisⱼ = basisⱼ