[GitHub] [spark] WeichenXu123 commented on a diff in pull request #39369: [SPARK-41775][PYTHON][ML] Adding support for PyForch functions

WeichenXu123 commented on code in PR #39369: URL: https://github.com/apache/spark/pull/39369#discussion_r1082450887 ## python/pyspark/ml/torch/tests/test_distributor.py: ## @@ -224,8 +293,10 @@ def setUp(self) -> None: self.sc = SparkContext("local-cluster[2,2,1024]",

[GitHub] [spark] WeichenXu123 commented on a diff in pull request #39369: [SPARK-41775][PYTHON][ML] Adding support for PyForch functions

WeichenXu123 commented on code in PR #39369: URL: https://github.com/apache/spark/pull/39369#discussion_r1082443370 ## python/pyspark/ml/torch/distributor.py: ## @@ -495,32 +546,119 @@ def set_gpus(context: "BarrierTaskContext") -> None: def _run_distributed_training(