code: def _train_logistic_regression_model_worker_fn(
      num_samples_per_worker: int,
      num_features: int,
      batch_size: int,
      max_iter: int,
      num_classes: int,
      learning_rate: float,
      momentum: float,
      fit_intercept: bool,
      seed: int,
  ) -> Any:
      from pyspark.ml.torch.distributor import (
          _get_spark_partition_data_loader,
      )
      from torch.nn.parallel import DistributedDataParallel as DDP
      import torch.distributed
      import torch.optim as optim
  
      # TODO: add a setting seed param.
      torch.manual_seed(seed)
  
      # TODO: support training on GPU
      # TODO: support L1 / L2 regularization
      torch.distributed.init_process_group("gloo")
  
      linear_model = torch_nn.Linear(
          num_features,
          num_classes,
          bias=fit_intercept,
          dtype=torch.float32,
      )
      ddp_model = DDP(linear_model)
  
      loss_fn = torch_nn.CrossEntropyLoss()
  
      optimizer = optim.SGD(
          ddp_model.parameters(), lr=learning_rate, momentum=momentum
      )
      data_loader = _get_spark_partition_data_loader(
          num_samples_per_worker,
          batch_size,
          num_workers=0,
          prefetch_factor=None,  # type: ignore
      )
      for i in range(max_iter):
          ddp_model.train()
  
          step_count = 0
  
          loss_sum = 0.0
          for x, target in data_loader:
              optimizer.zero_grad()
              output = ddp_model(x.to(torch.float32))
              loss = loss_fn(output, target.to(torch.long))
              loss.backward()
              loss_sum += loss.detach().numpy()
              optimizer.step()
              step_count += 1
  
          # TODO: early stopping
          #  When each epoch ends, computes loss on validation dataset and compare
          #  current epoch validation loss with last epoch validation loss, if
          #  less than provided `tol`, stop training.
  
          if torch.distributed.get_rank() == 0:
              print(
                  f"Progress: train epoch {i + 1} completes, train loss = {loss_sum / step_count}"
              )
  
      if torch.distributed.get_rank() == 0:
          return ddp_model.module.state_dict()
  
      return None
  
imports: from typing import Any, Dict, Union, List, Tuple, Callable, Optional
  import math
  import torch
  import torch.nn
  import numpy
  import pandas
  from pyspark import keyword_only
  from pyspark.ml.connect.base import _PredictorParams
  from pyspark.ml.param.shared import HasProbabilityCol
  from pyspark.sql import DataFrame
  from pyspark.ml.common import inherit_doc
  from pyspark.ml.torch.distributor import TorchDistributor
  from pyspark.ml.param.shared import HasMaxIter, HasFitIntercept, HasTol, HasWeightCol, HasSeed, HasNumTrainWorkers, HasBatchSize, HasLearningRate, HasMomentum
  from pyspark.ml.connect.base import Predictor, PredictionModel
  from pyspark.ml.connect.io_utils import ParamsReadWrite, CoreModelReadWrite
  from pyspark.sql import functions
label_desc: Use image node_id 1 for calling a global function with example usage: _train_logistic_regression_model_worker_fn(num_samples_per_worker, num_features, batch_size, max_iter, num_classes, learning_rate, momentum, fit_intercept, seed) and returns: None, ddp_model