Use correct split when doing reigen

1 files changed, 5 insertions, 5 deletions

diff --git a/train.py b/train.py
index 9bfef52..dd3633f 100755
--- a/train.py
+++ b/train.py
@@ -251,17 +251,17 @@ def main():
             target_pred, distances[i] = test_model(args.eigen, faces_train[i], faces_test, target_train[i], target_test, args)
         target_pred = np.argmin(distances, axis=0)
     elif args.reigen:
-        target_pred = np.zeros((args.reigen-args.eigen, 2*n_faces))
-        accuracy    = np.zeros((args.reigen-args.eigen, 2*n_faces))
-        rec_error   = np.zeros((args.reigen-args.eigen, 2*n_faces))
+        target_pred = np.zeros((args.reigen-args.eigen, target_test.shape[0]))
+        accuracy    = np.zeros(args.reigen-args.eigen)
+        rec_error   = np.zeros((args.reigen-args.eigen, target_test.shape[0]))
 
         for M in range(args.eigen, args.reigen):
             start = timer()
-            target_pred[i], rec_error[M - args.eigen] = test_model(M, faces_train, faces_test, target_train, target_test, args)
+            target_pred[M - args.eigen], rec_error[M - args.eigen] = test_model(M, faces_train, faces_test, target_train, target_test, args)
             end = timer()
             print("Run with", M, "eigenvalues completed in ", end-start, "seconds")
             print("Memory Used:", psutil.Process(os.getpid()).memory_info().rss)
-            accuracy[i] = accuracy_score(target_test, target_pred[i])
+            accuracy[M - args.eigen] = accuracy_score(target_test, target_pred[M-args.eigen])
         # Plot
         print('Max efficiency of ', max(accuracy), '% for M =', np.argmax(accuracy))
         plt.plot(range(args.eigen, args.reigen), 100*accuracy)