diff --git a/Homework1/conjugate_gradient.py b/Homework1/conjugate_gradient.py
index e2e54bb..cfe2d61 100644
--- a/Homework1/conjugate_gradient.py
+++ b/Homework1/conjugate_gradient.py
@@ -1,22 +1,26 @@
 import sys
 import numpy as np
 import argparse
 
 # Arguments parser
 my_parser = argparse.ArgumentParser()
 my_parser.add_argument('-b', action='store', type=float, nargs='+', required=True, help='The vector b in rowmajor format: numbers separated by space')
 my_parser.add_argument('-A', action='store', type=float, nargs='+', required=True, help='The matrix A in rowmajor format: numbers separated by space')
 
 args = my_parser.parse_args()
 
 # Extracting matrix A and vecor b from the input arguments
 A=np.array(args.A)
 b=np.array(args.b)
 
 # Checking the compatibility of size of A and b
 n=np.size(b)
-if np.size(A) != n:
+if np.size(A) != n*n:
     print('\n=====ERROR=====\n'
           'Matrix A and vector b are not compatible in size.\n'
           '===============\n')
-    exit()
\ No newline at end of file
+    exit()
+    
+A=A.reshape((n,n))
+
+print(A)
\ No newline at end of file