diff --git a/src/UniformTensorTrains/UniformTensorTrains.jl b/src/UniformTensorTrains/UniformTensorTrains.jl
index 3265e8c..8ce3f3d 100644
--- a/src/UniformTensorTrains/UniformTensorTrains.jl
+++ b/src/UniformTensorTrains/UniformTensorTrains.jl
@@ -5,6 +5,7 @@ using ..TensorTrains: _reshape1
 
 using LinearAlgebra
 using KrylovKit: eigsolve
+using TensorCast: TensorCast, @cast
 using Tullio: @tullio
 
 export AbstractPeriodicTensorTrain, PeriodicTensorTrain, flat_periodic_tt, rand_periodic_tt,
diff --git a/src/UniformTensorTrains/transfer_operator.jl b/src/UniformTensorTrains/transfer_operator.jl
index 3dbd351..02c0243 100644
--- a/src/UniformTensorTrains/transfer_operator.jl
+++ b/src/UniformTensorTrains/transfer_operator.jl
@@ -48,8 +48,7 @@ end
 
 function leading_eig(G::AbstractTransferOperator)
     GG = collect(G)
-    # @cast B[(i,j),(k,l)] := GG[i,j,k,l]
-    B = reshape(GG, prod(size(GG)[1:2]), prod(size(GG)[3:4]))
+    @cast B[(i,j),(k,l)] := GG[i,j,k,l]
     valsR, vecsR = eigsolve(B)
     valsL, vecsL = eigsolve(B')
     valsR[1] ≈ valsL[1] || @warn "Leading eigenvalue for A and Aᵀ not equal, got $(valsR[1]) and $(valsL[1])"