WeichenXu123 commented on code in PR #40607:
URL: https://github.com/apache/spark/pull/40607#discussion_r1155430451


##########
python/pyspark/ml/torch/distributor.py:
##########
@@ -554,7 +568,7 @@ def set_gpus(context: "BarrierTaskContext") -> None:
                     pass
 
             if context.partitionId() == 0:
-                yield output
+                yield pd.DataFrame(data={"output": 
[cloudpickle.dumps(output)]})

Review Comment:
   Q: Did you test what happens if the output exceeding 2GB ? I guess it is not 
supported. Shall we split it in the case ?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to