Add unthunk in tsvd!_pullback

Jutho · Jan 17, 2025 · 5a48e78 · 5a48e78
1 parent 4ef17d5
commit 5a48e78
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/ext/TensorKitChainRulesCoreExt/factorizations.jl b/ext/TensorKitChainRulesCoreExt/factorizations.jl
@@ -20,7 +20,8 @@ function ChainRulesCore.rrule(::typeof(TensorKit.tsvd!), t::AbstractTensorMap;
         Ũ, Σ̃, Ṽ⁺ = U, Σ, V⁺
     end
 
-    function tsvd!_pullback((ΔU, ΔΣ, ΔV⁺, Δϵ))
+    function tsvd!_pullback(ΔUSVϵ)
+        ΔU, ΔΣ, ΔV⁺, = unthunk.(ΔUSVϵ)
         Δt = similar(t)
         for (c, b) in blocks(Δt)
             Uc, Σc, V⁺c = block(U, c), block(Σ, c), block(V⁺, c)