update

stage_two.py

@@ -98,10 +98,10 @@ args = {
 task.connect(args)
 
 mnist_train = Dataset.get(
-    dataset_name=args['train_dataset_name'], dataset_project=args['dataset_project']
+    dataset_name=args['dataset_name_training'], dataset_project=args['dataset_project']
 ).get_local_copy()
 mnist_test = Dataset.get(
-    dataset_name=args['test_dataset_name'], dataset_project=args['dataset_project']
+    dataset_name=args['dataset_name_test'], dataset_project=args['dataset_project']
 ).get_local_copy()
 
 # get logger