diff --git a/image_ref/main.py b/image_ref/main.py
index ff5690919a3142d717eb3b3603f3987b3fd39c39..e2dbdeb98c21c6f6c38a122eb015183549eebec1 100644
--- a/image_ref/main.py
+++ b/image_ref/main.py
@@ -107,10 +107,10 @@ def run_duo(args):
     plt.ylim(0, 1.05)
     plt.show()
 
-    plt.savefig('output/training_plot_contrastive_noise_{}_lr_{}_model_{}_{}.png'.format(args.noise_threshold,args.lr,args.model,args.model_type))
+    plt.savefig('output/training_plot_contrastive_noise_{}_lr_{}_model_{}.png'.format(args.noise_threshold,args.lr,args.model))
     #load and evaluate best model
     load_model(model, args.save_path)
-    make_prediction_duo(model,data_test, 'output/confusion_matrix_contractive_noise_{}_lr_{}_model_{}_{}.png'.format(args.noise_threshold,args.lr,args.model,args.model_type))
+    make_prediction_duo(model,data_test, 'output/confusion_matrix_contractive_noise_{}_lr_{}_model_{}_{}.png'.format(args.noise_threshold,args.lr,args.model))
 
 
 def make_prediction_duo(model, data, f_name):