From f90f8f06e2847bc3109e76bf0b42e1545c29535b Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Fri, 8 May 2020 22:07:43 +0300
Subject: [PATCH] Fix Task.get_tasks() return type to Task

---
 trains/task.py | 13 ++++---------
 1 file changed, 4 insertions(+), 9 deletions(-)

diff --git a/trains/task.py b/trains/task.py
index 923624b0..98f0ce77 100644
--- a/trains/task.py
+++ b/trains/task.py
@@ -2144,16 +2144,11 @@ class Task(_Task):
         if task_ids:
             if isinstance(task_ids, six.string_types):
                 task_ids = [task_ids]
-            return [
-                cls(private=cls.__create_protection, task_id=task, log_to_backend=False)
-                for task in task_ids
-            ]
+            return [cls(private=cls.__create_protection, task_id=task_id, log_to_backend=False)
+                    for task_id in task_ids]
 
-        return cls._query_tasks(
-            project_name=project_name,
-            task_name=task_name,
-            **kwargs
-        )
+        return [cls(private=cls.__create_protection, task_id=task.id, log_to_backend=False)
+                for task in cls._query_tasks(project_name=project_name, task_name=task_name, **kwargs)]
 
     @classmethod
     def _query_tasks(cls, task_ids=None, project_name=None, task_name=None, **kwargs):