Add some changes and result

src/models/train_model.py

@@ -63,6 +63,7 @@ def parser_process(value, state):
 " Code copied from: https://stackoverflow.com/questions/48391777/nargs-equivalent-for-options-in-click"
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+device = torch.device('cuda:0')
 
 pd.options.mode.chained_assignment = None  # default='warn'