From 938012ac77a19e564ad968ad8dfee61d818d4c7d Mon Sep 17 00:00:00 2001
From: Joeper214 <joeper.serrano@gmail.com>
Date: Thu, 20 Feb 2020 13:30:24 +0900
Subject: [PATCH] Use list instead

---
 blueoil/cmd/init.py | 15 +++++++--------
 1 file changed, 7 insertions(+), 8 deletions(-)

diff --git a/blueoil/cmd/init.py b/blueoil/cmd/init.py
index b71f4b955..f25e3ec1b 100644
--- a/blueoil/cmd/init.py
+++ b/blueoil/cmd/init.py
@@ -160,14 +160,13 @@ def dataset_format_choices(task_type):
 
 
 def default_batch_size(task_type):
-    if task_type == 'classification':
-        return '10'
-    elif task_type == 'object_detection':
-        return '16'
-    elif task_type == 'semantic_segmentation':
-        return '8'
-    elif task_type == 'keypoint_detection':
-        return '4'
+    default_task_type_batch_sizes = {
+        'classification': '10',
+        'object_detection': '16',
+        'semantic_segmentation': '8',
+        'keypoint_detection': '4',
+    }
+    return default_task_type_batch_sizes[task_type]
 
 
 def prompt(question):