From b6306c8eefd6262ecfc4a00bc5fefed36fbc19c8 Mon Sep 17 00:00:00 2001
From: Theodore Vasiloudis <thvasilo@amazon.com>
Date: Thu, 19 Dec 2024 21:14:40 +0000
Subject: [PATCH] Fix docstring

---
 .../data_transformations/dist_label_loader.py                 | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/graphstorm-processing/graphstorm_processing/data_transformations/dist_label_loader.py b/graphstorm-processing/graphstorm_processing/data_transformations/dist_label_loader.py
index cd0ba563f..a85299947 100644
--- a/graphstorm-processing/graphstorm_processing/data_transformations/dist_label_loader.py
+++ b/graphstorm-processing/graphstorm_processing/data_transformations/dist_label_loader.py
@@ -75,9 +75,9 @@ class CustomSplit:
     ----------
     train : list[str]
         Paths of the training mask parquet files.
-    valid : str
+    valid : list[str]
         Paths of the validation mask parquet files.
-    test : str
+    test : list[str]
         Paths of the testing mask parquet files.
     mask_columns : list[str]
         List of columns that contain original string ids.