config solo or duo

config/config.py

@@ -11,6 +11,7 @@ def load_args():
     parser.add_argument('--lr', type=float, default=0.001)
     parser.add_argument('--batch_size', type=int, default=64)
     parser.add_argument('--model', type=str, default='ResNet18')
+    parser.add_argument('--model_type', type=str, default='solo')
     parser.add_argument('--dataset_dir', type=str, default='data/processed_data/png_image/data_training')
     parser.add_argument('--output', type=str, default='output/out.csv')
     parser.add_argument('--save_path', type=str, default='output/best_model.pt')

main.py

@@ -206,10 +206,10 @@ def run_duo(args):
     plt.plot(train_acc)
     plt.ylim(0, 1.05)
     plt.show()
-    plt.savefig('output/training_plot_noise_{}_lr_{}_model_{}.png'.format(args.noise_threshold,args.lr,args.model))
+    plt.savefig('output/training_plot_noise_{}_lr_{}_model_{}_{}.png'.format(args.noise_threshold,args.lr,args.model,args.model_type))
 
     load_model(model, args.save_path)
-    make_prediction_duo(model,data_test, 'output/confusion_matrix_noise_{}_lr_{}_model_{}.png'.format(args.noise_threshold,args.lr,args.model))
+    make_prediction_duo(model,data_test, 'output/confusion_matrix_noise_{}_lr_{}_model_{}_{}.png'.format(args.noise_threshold,args.lr,args.model,args.model_type))
 
 
 def make_prediction_duo(model, data, f_name):
@@ -254,4 +254,7 @@ def load_model(model, path):
 
 if __name__ == '__main__':
     args = load_args()
-    run(args)
\ No newline at end of file
+    if args.model_type=='duo':
+        run_duo(args)
+    else :
+        run(args)
\ No newline at end of file