diff --git a/scratchpad/qtensor_MPS/mpo.py b/scratchpad/qtensor_MPS/mpo.py
index e2d746e3..b70bda3c 100644
--- a/scratchpad/qtensor_MPS/mpo.py
+++ b/scratchpad/qtensor_MPS/mpo.py
@@ -62,20 +62,18 @@ def get_mpo_nodes(self, original) -> list[tn.Node]:
     def get_mpo_node(self, index, original) -> list[tn.Node]:
         return self.get_mpo_nodes(original)[index]
 
-    def construct_mpo(self, gate_function, tensor_name, N, physical_dim=1) -> "MPO":
+    def construct_mpo(self, gate_function) -> "MPO":
         # IIZ
-        gate_function = gate_function.reshape([physical_dim] * 2 * N)
-        print(gate_function.shape)
+        gate_function = gate_function.reshape([self.physical_dim] * 2 * self.N)
         to_split = tn.Node(
             gate_function,
-            axis_names=["u" + str(i) for i in range(N)]
-            + ["d" + str(i) for i in range(N)],
+            axis_names=["u" + str(i) for i in range(self.N)]
+            + ["d" + str(i) for i in range(self.N)],
         )
-        print(to_split.shape)
 
         nodes = []
 
-        for i in range(N - 1):
+        for i in range(self.N - 1):
             left_edges = []
             right_edges = []
 
@@ -94,12 +92,12 @@ def construct_mpo(self, gate_function, tensor_name, N, physical_dim=1) -> "MPO":
                 to_split,
                 left_edges,
                 right_edges,
-                left_name=tensor_name + str(i),
+                left_name=self.tensor_name + str(i),
                 max_singular_values=1,
             )
             nodes.append(left)
             to_split = right
-        to_split.name = tensor_name + str(N - 1)
+        to_split.name = self.tensor_name + str(self.N - 1)
         nodes.append(to_split)
 
         self._nodes = nodes
diff --git a/scratchpad/qtensor_MPS/test.py b/scratchpad/qtensor_MPS/test.py
index c88cf459..c95cb81a 100644
--- a/scratchpad/qtensor_MPS/test.py
+++ b/scratchpad/qtensor_MPS/test.py
@@ -192,25 +192,50 @@ def test_mpo_two_qubit():
     # assert (mps1.get_wavefunction() == mps2.get_wavefunction()).all()
 
 
-def test_mpo_construction_from_gate_function():
-    # IZI
+def test_mpo_construction_from_single_qubit_gate_function():
+    # IX
     I = np.eye(2)
-    Z = np.array([[0.0, 1.0], [1.0, 0.0]], dtype=np.complex64)
-    H = 1 / np.sqrt(2) * np.array([[1.0, 1.0], [1.0, -1.0]], dtype=np.complex64)
-    gate_func = np.kron(I, I)
-    print(gate_func.shape)
+    X = xmatrix
+
+    mps2 = MPS("q", 2, 2)
+    mps2.apply_single_qubit_gate(xgate(), 0)
+    print("MPS with gate", mps2.get_wavefunction())
+
+    gate_func = np.kron(X, I)
+    mpo = MPOLayer("q", 2, 2)
+    mpo.construct_mpo(gate_func)
+    mps = MPS("q", 2, 2)
+
+    mps.apply_mpo_layer(mpo)
+    print("MPS with MPO", mps.get_wavefunction())
+
+    condition = np.allclose(
+        np.array(mps.get_wavefunction()),
+        np.array(mps2.get_wavefunction()),
+        rtol=1e-05,
+        atol=1e-08,
+    )
+    message = "Arrays are not equal within tolerance"
+
+    assert condition, message
+
+
+def test_mpo_construction_from_two_qubit_gate_function():
+    I = np.eye(2)
+    gate_func = cnot_matrix
     mpo = MPOLayer("q", 2, 2)
-    mpo.construct_mpo(gate_func, "q", 2, 2)
+    mpo.construct_mpo(gate_func)
     mps = MPS("q", 2, 2)
-    print(mps.get_wavefunction())
+    print("INITIAL MPS", mps.get_wavefunction())
     mps.apply_mpo_layer(mpo)
-    print(mps.get_wavefunction())
+    print("FINAL MPS", mps.get_wavefunction())
 
 
 # TODO:
 # Verify for two qubit gates. - (SVD and then apply it) (apply and then SVD)
-# Evolve half od circuit in mps and half in mpo.
+# Evolve half of circuit in mps and half in mpo.
 # Evaluate Bell state, epectation value of X at a particular index
 # Evaluate X, Y, Z observable
 # Check for one/two/three qubits (cover edge cases) for both mps/mpo
-test_mpo_construction_from_gate_function()
+# test_mpo_construction_from_single_qubit_gate_function()
+# test_mpo_construction_from_two_qubit_gate_function()