From 99a8724bece05c5794d98a1cee584573e82b9e05 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Tue, 1 Sep 2020 17:46:24 +0300
Subject: [PATCH] Fix Task.wait_for_status() to reload after it ends

---
 trains/task.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/trains/task.py b/trains/task.py
index fc001b99..92c11ec4 100644
--- a/trains/task.py
+++ b/trains/task.py
@@ -1600,6 +1600,9 @@ class Task(_Task):
         if raise_on_status and self.status in raise_on_status:
             raise RuntimeError("Task {} has status: {}.".format(self.task_id, self.status))
 
+        # make sure we have the Task object
+        self.reload()
+
     def export_task(self):
         # type: () -> dict
         """