diff --git a/python/gigl/utils/data_splitters.py b/python/gigl/utils/data_splitters.py
index a7272cd4d..259e0b047 100644
--- a/python/gigl/utils/data_splitters.py
+++ b/python/gigl/utils/data_splitters.py
@@ -684,7 +684,7 @@ def _get_padded_labels(
         ],
         dim=0,
     )
-    return labels
+    return labels[:,:3]
 
 
 def _assert_sampling_direction(sampling_direction: str):