Update train_model.py

src/models/train_model.py

@@ -30,6 +30,9 @@ def train_model():
 
     model.save_model(model_dir=params['model_dir'])
 
+    with open('wandb/latest-run/files/wandb-summary.json') as json_file:
+        data = json.load(json_file)
+
 
 if __name__ == '__main__':
     train_model()