diff --git a/python/tasks/jigsaw_toxic.py b/python/tasks/jigsaw_toxic.py
index 51513b4..f1d0749 100644
--- a/python/tasks/jigsaw_toxic.py
+++ b/python/tasks/jigsaw_toxic.py
@@ -229,6 +229,10 @@ def kernel_2():
         xtrain_seq=xtrain_seq,
         token=token,
         max_len=max_len,
+        xtrain=xtrain,
+        xvalid=xvalid,
+        ytrain=ytrain,
+        yvalid=yvalid,
     )