[SPARK-53453][PYTHON][ML] Unblock 'torch<2.6.0'

### What changes were proposed in this pull request?
Unblock 'torch<2.6.0'

### Why are the changes needed?
to test with latest torch

### Does this PR introduce _any_ user-facing change?
No

### How was this patch tested?
CI

### Was this patch authored or co-authored using generative AI tooling?
No

Closes #52197 from zhengruifeng/torch_280.

Authored-by: Ruifeng Zheng <ruifengz@apache.org>
Signed-off-by: Dongjoon Hyun <dongjoon@apache.org>

Author: zhengruifeng

dev/spark-test-image/python-310/Dockerfile

@@ -72,6 +72,6 @@ RUN curl -sS https://bootstrap.pypa.io/get-pip.py | python3.10
 RUN python3.10 -m pip install --ignore-installed blinker>=1.6.2 # mlflow needs this
 RUN python3.10 -m pip install --ignore-installed 'six==1.16.0'  # Avoid `python3-six` installation
 RUN python3.10 -m pip install $BASIC_PIP_PKGS unittest-xml-reporting $CONNECT_PIP_PKGS && \
-    python3.10 -m pip install 'torch<2.6.0' torchvision --index-url https://download.pytorch.org/whl/cpu && \
+    python3.10 -m pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu && \
     python3.10 -m pip install deepspeed torcheval && \
     python3.10 -m pip cache purge

dev/spark-test-image/python-311-classic-only/Dockerfile

@@ -74,6 +74,6 @@ ARG TEST_PIP_PKGS="coverage unittest-xml-reporting"
 RUN curl -sS https://bootstrap.pypa.io/get-pip.py | python3.11
 RUN python3.11 -m pip install --ignore-installed blinker>=1.6.2 # mlflow needs this
 RUN python3.11 -m pip install $BASIC_PIP_PKGS $TEST_PIP_PKGS && \
-    python3.11 -m pip install 'torch<2.6.0' torchvision --index-url https://download.pytorch.org/whl/cpu && \
+    python3.11 -m pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu && \
     python3.11 -m pip install deepspeed torcheval && \
     python3.11 -m pip cache purge

dev/spark-test-image/python-311/Dockerfile

@@ -75,6 +75,6 @@ ARG CONNECT_PIP_PKGS="grpcio==1.67.0 grpcio-status==1.67.0 protobuf==5.29.1 goog
 RUN curl -sS https://bootstrap.pypa.io/get-pip.py | python3.11
 RUN python3.11 -m pip install --ignore-installed blinker>=1.6.2 # mlflow needs this
 RUN python3.11 -m pip install $BASIC_PIP_PKGS unittest-xml-reporting $CONNECT_PIP_PKGS && \
-    python3.11 -m pip install 'torch<2.6.0' torchvision --index-url https://download.pytorch.org/whl/cpu && \
+    python3.11 -m pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu && \
     python3.11 -m pip install deepspeed torcheval && \
     python3.11 -m pip cache purge

dev/spark-test-image/python-312/Dockerfile

@@ -75,6 +75,6 @@ ARG CONNECT_PIP_PKGS="grpcio==1.67.0 grpcio-status==1.67.0 protobuf==5.29.1 goog
 RUN curl -sS https://bootstrap.pypa.io/get-pip.py | python3.12
 RUN python3.12 -m pip install --ignore-installed blinker>=1.6.2 # mlflow needs this
 RUN python3.12 -m pip install $BASIC_PIP_PKGS unittest-xml-reporting $CONNECT_PIP_PKGS lxml && \
-    python3.12 -m pip install 'torch<2.6.0' torchvision --index-url https://download.pytorch.org/whl/cpu && \
+    python3.12 -m pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu && \
     python3.12 -m pip install torcheval && \
     python3.12 -m pip cache purge

dev/spark-test-image/python-313/Dockerfile

@@ -75,6 +75,6 @@ ARG CONNECT_PIP_PKGS="grpcio==1.67.0 grpcio-status==1.67.0 protobuf==5.29.1 goog
 RUN curl -sS https://bootstrap.pypa.io/get-pip.py | python3.13
 RUN python3.13 -m pip install --ignore-installed blinker>=1.6.2 # mlflow needs this
 RUN python3.13 -m pip install $BASIC_PIP_PKGS unittest-xml-reporting $CONNECT_PIP_PKGS lxml && \
-    python3.13 -m pip install 'torch<2.6.0' torchvision --index-url https://download.pytorch.org/whl/cpu && \
+    python3.13 -m pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu && \
     python3.13 -m pip install torcheval && \
     python3.13 -m pip cache purge

python/pyspark/ml/connect/classification.py

@@ -381,7 +381,7 @@ def _save_core_model(self, path: str) -> None:
     def _load_core_model(self, path: str) -> None:
         import torch
 
-        lor_torch_model = torch.load(path)
+        lor_torch_model = torch.load(path, weights_only=False)
         self.torch_model = lor_torch_model[0]
 
     def _get_extra_metadata(self) -> Dict[str, Any]:

python/pyspark/ml/tests/connect/test_legacy_mode_classification.py

@@ -183,7 +183,8 @@ def test_save_load(self):
 
             # test saved torch model can be loaded by pytorch solely
             lor_torch_model = torch.load(
-                os.path.join(local_model_path, "LogisticRegressionModel.torch")
+                os.path.join(local_model_path, "LogisticRegressionModel.torch"),
+                weights_only=False,
             )
 
             with torch.inference_mode():