[dask] Allow empty data matrix in AFT survival

dmlc · Nov 12, 2020 · eb58239 · eb58239
1 parent d711d64
commit eb58239
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 6 deletions.
diff --git a/src/metric/survival_metric.cu b/src/metric/survival_metric.cu
@@ -206,10 +206,6 @@ struct EvalEWiseSurvivalBase : public Metric {
   bst_float Eval(const HostDeviceVector<bst_float>& preds,
                  const MetaInfo& info,
                  bool distributed) override {
-    CHECK_NE(info.labels_lower_bound_.Size(), 0U)
-      << "labels_lower_bound cannot be empty";
-    CHECK_NE(info.labels_upper_bound_.Size(), 0U)
-      << "labels_upper_bound cannot be empty";
     CHECK_EQ(preds.Size(), info.labels_lower_bound_.Size());
     CHECK_EQ(preds.Size(), info.labels_upper_bound_.Size());
 

diff --git a/tests/python/test_with_dask.py b/tests/python/test_with_dask.py
@@ -594,7 +594,6 @@ def test_predict_with_meta(client):
 
 
 def run_aft_survival(client, dmatrix_t):
-    # survival doesn't handle empty dataset well.
     df = dd.read_csv(os.path.join(tm.PROJECT_ROOT, 'demo', 'data',
                                   'veterans_lung_cancer.csv'))
     y_lower_bound = df['Survival_label_lower_bound']
@@ -632,7 +631,7 @@ def run_aft_survival(client, dmatrix_t):
 
 
 def test_aft_survival():
-    with LocalCluster(n_workers=1) as cluster:
+    with LocalCluster(n_workers=kWorkers) as cluster:
         with Client(cluster) as client:
             run_aft_survival(client, DaskDMatrix)