diff --git a/Hamiltonian.jl b/Hamiltonian.jl
index 19aabfa..9a6fd73 100644
--- a/Hamiltonian.jl
+++ b/Hamiltonian.jl
@@ -33,7 +33,7 @@ struct Hamiltonian{T}
     end
 end
 
-Base.size(H::Hamiltonian, i::Int)::Int = (i == 1 || i == 2) ? H.s.N^(H.s.d * (H.s.n - 2)) * H.dim1 : throw(ArgumentError("Hamiltonian only has 2 dimesions"))
+Base.size(H::Hamiltonian, i::Int)::Int = (i == 1 || i == 2) ? H.s.N^(H.s.d * (H.s.n - 2)) * length(H.s.unique_i) : throw(ArgumentError("Hamiltonian only has 2 dimesions"))
 Base.size(H::Hamiltonian)::Dims{2} = (size(H, 1), size(H, 2))
 
 "Dimensions of a vector to which 'H' can be applied"