Fix test.

dmlc · Feb 8, 2021 · 8802c17 · 8802c17
1 parent 34cc13e
commit 8802c17
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/tests/python-gpu/test_gpu_with_dask.py b/tests/python-gpu/test_gpu_with_dask.py
@@ -304,9 +304,15 @@ def worker_fn(worker_addr: str, data_ref: Dict) -> None:
 
             futures = []
             for i in range(len(workers)):
-                futures.append(client.submit(worker_fn, workers[i],
-                                             m.create_fn_args(workers[i]), pure=False,
-                                             workers=[workers[i]]))
+                futures.append(
+                    client.submit(
+                        worker_fn,
+                        workers[i],
+                        m._create_fn_args(workers[i]),
+                        pure=False,
+                        workers=[workers[i]]
+                    )
+                )
             client.gather(futures)
 
     def test_interface_consistency(self) -> None: