diff --git a/train.py b/train.py
index 78e5b1f..10893f7 100644
--- a/train.py
+++ b/train.py
@@ -46,7 +46,12 @@ def train(opt):
         shutil.rmtree(opt.log_path)
     os.makedirs(opt.log_path)
     writer = SummaryWriter(opt.log_path)
-    optimizer = torch.optim.Adam(model.parameters(), lr=1e-6)
+    
+    if opt.optimizer == "adam":    
+        optimizer = torch.optim.Adam(model.parameters(), opt.lr)
+    else:
+        optimizer = torch.optim.SGD(model.parameters(), opt.lr)
+    
     criterion = nn.MSELoss()
     game_state = FlappyBird()
     image, reward, terminal = game_state.next_frame(0)
@@ -71,7 +76,7 @@ def train(opt):
             action = randint(0, 1)
         else:
 
-            action = torch.argmax(prediction)[0]
+            action = torch.argmax(prediction)
 
         next_image, reward, terminal = game_state.next_frame(action)
         next_image = pre_processing(next_image[:game_state.screen_width, :int(game_state.base_y)], opt.image_size,