diff --git a/Manifest.toml b/Manifest.toml
index a9e1a40..f7271f0 100644
--- a/Manifest.toml
+++ b/Manifest.toml
@@ -1,5 +1,17 @@
 # This file is machine-generated - editing it directly is not advised
 
+[[AbstractFFTs]]
+deps = ["LinearAlgebra"]
+git-tree-sha1 = "380e36c66edfa099cd90116b24c1ce8cafccac40"
+uuid = "621f4979-c628-5d54-868e-fcf4e3e8185c"
+version = "0.4.1"
+
+[[Adapt]]
+deps = ["LinearAlgebra"]
+git-tree-sha1 = "82dab828020b872fa9efd3abec1152b075bc7cbf"
+uuid = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+version = "1.0.0"
+
 [[ArrayInterface]]
 deps = ["LinearAlgebra", "Requires", "SparseArrays"]
 git-tree-sha1 = "981354dab938901c2b607a213e62d9defa50b698"
@@ -35,9 +47,9 @@ version = "0.5.6"
 
 [[CSTParser]]
 deps = ["Tokenize"]
-git-tree-sha1 = "0ff80f68f55fcde2ed98d7b24d7abaf20727f3f8"
+git-tree-sha1 = "c69698c3d4a7255bc1b4bc2afc09f59db910243b"
 uuid = "00ebfdb7-1f24-5e51-bd34-a7502290713f"
-version = "0.6.1"
+version = "0.6.2"
 
 [[Calculus]]
 deps = ["Compat"]
@@ -63,6 +75,12 @@ git-tree-sha1 = "84aa74986c5b9b898b0d1acaf3258741ee64754f"
 uuid = "34da2185-b29b-5c13-b0c7-acf172513d20"
 version = "2.1.0"
 
+[[Conda]]
+deps = ["JSON", "VersionParsing"]
+git-tree-sha1 = "9a11d428dcdc425072af4aea19ab1e8c3e01c032"
+uuid = "8f4d0f93-b110-5947-807f-2305c1781a2d"
+version = "1.3.0"
+
 [[Crayons]]
 deps = ["Test"]
 git-tree-sha1 = "f621b8ef51fd2004c7cf157ea47f027fdeac5523"
@@ -70,9 +88,9 @@ uuid = "a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f"
 version = "4.0.0"
 
 [[DataAPI]]
-git-tree-sha1 = "891a09f4f90361a28d0391c104a65c0202e22624"
+git-tree-sha1 = "8903f0219d3472543fc4b2f5ebaf675a07f817c0"
 uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
-version = "1.0.0"
+version = "1.0.1"
 
 [[DataStructures]]
 deps = ["InteractiveUtils", "OrderedCollections"]
@@ -89,10 +107,10 @@ deps = ["Mmap"]
 uuid = "8bb1440f-4735-579b-a4ab-409b98df4dab"
 
 [[DiffEqDiffTools]]
-deps = ["ArrayInterface", "LinearAlgebra", "SparseArrays", "StaticArrays"]
-git-tree-sha1 = "d6582a0e9e26a6453a4c03f2837b1049f7007a3f"
+deps = ["ArrayInterface", "LinearAlgebra", "Requires", "SparseArrays", "StaticArrays"]
+git-tree-sha1 = "21b855cb29ec4594f9651e0e9bdc0cdcfdcd52c1"
 uuid = "01453d9d-ee7c-5054-8395-0335cb756afa"
-version = "1.2.0"
+version = "1.3.0"
 
 [[DiffResults]]
 deps = ["Compat", "StaticArrays"]
@@ -107,9 +125,15 @@ uuid = "b552c78f-8df3-52c6-915a-8e097449b14b"
 version = "0.0.10"
 
 [[Distributed]]
-deps = ["LinearAlgebra", "Random", "Serialization", "Sockets"]
+deps = ["Random", "Serialization", "Sockets"]
 uuid = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 
+[[FFTW]]
+deps = ["AbstractFFTs", "BinaryProvider", "Conda", "Libdl", "LinearAlgebra", "Reexport", "Test"]
+git-tree-sha1 = "e1a479d3c972f20c9a70563eec740bbfc786f515"
+uuid = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
+version = "0.3.0"
+
 [[FillArrays]]
 deps = ["LinearAlgebra", "Random", "SparseArrays", "Test"]
 git-tree-sha1 = "9ab8f76758cbabba8d7f103c51dce7f73fcf8e92"
@@ -122,6 +146,12 @@ git-tree-sha1 = "4c4d727f1b7e0092134fabfab6396b8945c1ea5b"
 uuid = "f6369f11-7733-5829-9624-2563aa707210"
 version = "0.10.3"
 
+[[GPUArrays]]
+deps = ["Adapt", "FFTW", "FillArrays", "LinearAlgebra", "Printf", "Random", "Serialization", "StaticArrays", "Test"]
+git-tree-sha1 = "4310da9bc8d51dccbffbd77138546208d0da93bf"
+uuid = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
+version = "1.0.0"
+
 [[IRTools]]
 deps = ["InteractiveUtils", "MacroTools", "Test"]
 git-tree-sha1 = "a9b1fc7745ae4745a634bbb6d1cb7fd64e37248a"
@@ -129,7 +159,7 @@ uuid = "7869d1d1-7146-5819-86e3-90919afe41df"
 version = "0.2.2"
 
 [[InteractiveUtils]]
-deps = ["LinearAlgebra", "Markdown"]
+deps = ["Markdown"]
 uuid = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 
 [[IterTools]]
@@ -143,6 +173,11 @@ git-tree-sha1 = "b34d7cef7b337321e97d22242c3c2b91f476748e"
 uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
 version = "0.21.0"
 
+[[LRUCache]]
+git-tree-sha1 = "2d7617aad2896b4df54fffd868976f708ef16614"
+uuid = "8ac3fa9e-de4c-5943-b1dc-09c6b5f20637"
+version = "1.0.1"
+
 [[LibGit2]]
 uuid = "76f85450-5226-5b5a-8eaa-529ad045b433"
 
@@ -200,8 +235,10 @@ uuid = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
 version = "0.3.2"
 
 [[OMEinsum]]
-deps = ["Combinatorics", "IRTools", "LinearAlgebra", "MacroTools", "PkgBenchmark", "Requires", "TensorOperations", "TupleTools", "Zygote"]
-git-tree-sha1 = "385e12265ef737c7c2386dd27a546ed2733e1dfe"
+deps = ["Combinatorics", "GPUArrays", "IRTools", "LinearAlgebra", "MacroTools", "PkgBenchmark", "Requires", "TensorOperations", "Test", "TupleTools", "Zygote"]
+git-tree-sha1 = "83c50b3fb529936fff187df3f9011b4aca7a477a"
+repo-rev = "feat-fastergpu-notype"
+repo-url = "https://github.com/under-Peter/OMEinsum.jl.git"
 uuid = "ebe7aa44-baf0-506c-a96f-8464559b3922"
 version = "0.2.0"
 
@@ -269,6 +306,12 @@ uuid = "3fa0cd96-eef1-5676-8a61-b3b8758bbffb"
 deps = ["Serialization"]
 uuid = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 
+[[Reexport]]
+deps = ["Pkg"]
+git-tree-sha1 = "7b1d07f411bc8ddb7977ec7f377b97b158514fe0"
+uuid = "189a3867-3050-52da-a836-e630ba90ab69"
+version = "0.2.0"
+
 [[Requires]]
 deps = ["Test"]
 git-tree-sha1 = "f6fbf4ba64d295e146e49e021207993b6b48c7d1"
@@ -327,10 +370,10 @@ uuid = "5e0ebb24-38b0-5f93-81fe-25c709ecae67"
 version = "0.3.1"
 
 [[TensorOperations]]
-deps = ["LinearAlgebra", "Strided", "TupleTools"]
-git-tree-sha1 = "2804781ab5e899a0e4c61a4d655e4dea88553710"
+deps = ["LRUCache", "LinearAlgebra", "Strided", "TupleTools"]
+git-tree-sha1 = "15d31a5821ca106d567eb148df123ca57d8cf4c9"
 uuid = "6aa20fa7-93e2-5fca-9bc0-fbd0db3c71a2"
-version = "1.1.0"
+version = "1.2.0"
 
 [[Test]]
 deps = ["Distributed", "InteractiveUtils", "Logging", "Random"]
@@ -360,12 +403,18 @@ uuid = "30578b45-9adc-5946-b283-645ec420af67"
 version = "0.4.0"
 
 [[UUIDs]]
-deps = ["Random"]
+deps = ["Random", "SHA"]
 uuid = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
 
 [[Unicode]]
 uuid = "4ec0a83e-493e-50e2-b9ac-8f72acf5a8f5"
 
+[[VersionParsing]]
+deps = ["Compat"]
+git-tree-sha1 = "c9d5aa108588b978bd859554660c8a5c4f2f7669"
+uuid = "81def892-9a0e-5fdd-b105-ffc91e053289"
+version = "1.1.3"
+
 [[Zygote]]
 deps = ["DiffRules", "FillArrays", "ForwardDiff", "IRTools", "InteractiveUtils", "LinearAlgebra", "MacroTools", "NNlib", "NaNMath", "Random", "Requires", "SpecialFunctions", "Statistics"]
 git-tree-sha1 = "520d65c5e5554473863d738bde053f5f6769d3be"
diff --git a/src/ctmrg.jl b/src/ctmrg.jl
index 4ebb945..85896fa 100644
--- a/src/ctmrg.jl
+++ b/src/ctmrg.jl
@@ -6,11 +6,13 @@ using Base.Iterators: take, drop
 using Optim, LineSearches
 
 function initializec(a, χ, randinit)
-    c = zeros(eltype(a), χ, χ)
+    # c = zeros(eltype(a), χ, χ)
+    c = similar(a, χ, χ)
     if randinit
         rand!(c)
         c += adjoint(c)
     else
+        fill!(c, zero(eltype(a)))
         cinit = ein"ijkl -> ij"(a)
         foreach(CartesianIndices(cinit)) do i
             i in CartesianIndices(c) && (c[i] = cinit[i])
@@ -20,12 +22,13 @@ function initializec(a, χ, randinit)
 end
 
 function initializet(a, χ, randinit)
-    t = zeros(eltype(a), χ, size(a,1), χ)
+    t = similar(a, χ, size(a,1), χ)
     if randinit
         rand!(t)
         t += permutedims(conj(t), (3,2,1))
     else
         tinit = ein"ijkl -> ijk"(a)
+        fill!(t, zero(eltype(a)))
         foreach(CartesianIndices(tinit)) do i
             i in CartesianIndices(t) && (t[i] = tinit[i])
         end
@@ -40,16 +43,16 @@ function ctmrg(a, χ, tol, maxit::Integer, randinit = false)
     # initialize
     cinit = initializec(a, χ, randinit)
     tinit = initializet(a, χ, randinit)
-    oldvals = fill(Inf, χ*d)
-
+    oldvals = similar(real(a), χ*d) .* convert(real(eltype(a)),Inf)
     stopfun = StopFunction(oldvals, -1, tol, maxit)
+
     c, t, = fixedpoint(ctmrgstep, (cinit, tinit, oldvals), (a, χ, d), stopfun)
     return c, t
 end
 
 function ctmrgstep((c,t,vals), (a, χ, d))
     # grow
-    cp = ein"ad,iba,dcl,jkcb -> ijlk"(c, t, t, a)
+    cp = ein"((ad,iba),dcl),jkcb -> ijlk"(c, t, t, a)
     tp = ein"iam,jkla -> ijklm"(t,a)
 
     # renormalize
@@ -58,8 +61,8 @@ function ctmrgstep((c,t,vals), (a, χ, d))
     u, s, v = svd(cpmat)
     z = reshape(u[:, 1:χ], χ, d, χ)
 
-    c = ein"abcd,abi,cdj -> ij"(cp, conj(z), z)
-    t = ein"abjcd,abi,dck -> ijk"(tp, conj(z), z)
+    c = ein"(abcd,abi),cdj -> ij"(cp, conj(z), z)
+    t = ein"(abjcd,abi),dck -> ijk"(tp, conj(z), z)
 
     vals = s ./ s[1]
 
diff --git a/src/trg.jl b/src/trg.jl
index 35e4bf4..599968d 100644
--- a/src/trg.jl
+++ b/src/trg.jl
@@ -24,9 +24,11 @@ function trg(a::AbstractArray{T,4}, χ, niter; tol::Float64 = 1e-16) where T
         dr, ul = trg_svd(dr_ul, χ, tol)
         ld, ru = trg_svd(ld_ru, χ, tol)
 
-        a = ein"npu,por,dom,lmn -> urdl"(dr,ld,ul,ru)
+        a = ein"((npu,por),dom),lmn -> urdl"(dr,ld,ul,ru)
     end
-    trace = ein"ijij -> "(a)[]
+    # trace = ein"ijij -> "(a)[]
+    d1,d2,d3,d4 = size(a)
+    trace = ein"ii -> "(reshape(a,d1*d2, d3*d4))[]
     lnZ += log(trace)/2.0^niter
     return lnZ
 end