From 91d95c4b618679746ab96627431fa70f021ed92d Mon Sep 17 00:00:00 2001
From: Xin Yang <xyang19@gmail.com>
Date: Tue, 12 Nov 2024 08:57:32 -0800
Subject: [PATCH] [pytorch] Updates PyTorch to 2.5.1

---
 engines/pytorch/pytorch-jni/build.gradle.kts | 6 ++----
 gradle/libs.versions.toml                    | 2 +-
 2 files changed, 3 insertions(+), 5 deletions(-)

diff --git a/engines/pytorch/pytorch-jni/build.gradle.kts b/engines/pytorch/pytorch-jni/build.gradle.kts
index 43c967298304..b6b6342885db 100644
--- a/engines/pytorch/pytorch-jni/build.gradle.kts
+++ b/engines/pytorch/pytorch-jni/build.gradle.kts
@@ -1,5 +1,3 @@
-import java.net.URL
-
 plugins {
     ai.djl.javaProject
     ai.djl.publish
@@ -32,7 +30,7 @@ tasks {
                 "osx-aarch64/cpu/libdjl_torch.dylib",
                 "win-x86_64/cpu/djl_torch.dll"
             ) + when {
-                ptVersion.startsWith("2.4.") -> listOf(
+                ptVersion.matches(Regex("2.[4-5].\\d")) -> listOf(
                     "linux-x86_64/cu124/libdjl_torch.so",
                     "linux-x86_64/cu124-precxx11/libdjl_torch.so",
                     "win-x86_64/cu124/djl_torch.dll"
@@ -103,4 +101,4 @@ publishing {
             }
         }
     }
-}
\ No newline at end of file
+}
diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml
index 990bad2da5aa..2bbd623b9a54 100644
--- a/gradle/libs.versions.toml
+++ b/gradle/libs.versions.toml
@@ -3,7 +3,7 @@ format.version = "1.1"
 
 [versions]
 djl = "0.31.0"
-pytorch = "2.4.0"
+pytorch = "2.5.1"
 tensorflow = "2.16.1"
 tensorflowCore = "1.0.0-rc.1"
 mxnet = "1.9.1"