From f1b9da67e035ab0642681430fa898dfb9712d79c Mon Sep 17 00:00:00 2001
From: Frank Liu <frankfliu2000@gmail.com>
Date: Tue, 9 Jan 2024 13:43:46 -0800
Subject: [PATCH] [pytorch] Allows to load libstdc++.so.6 form different
 location (#2929)

---
 .../src/main/java/ai/djl/pytorch/jni/LibUtils.java    | 11 +++++++++++
 .../{LibUtilsTest.java => ALibUtilsTest.java}         | 10 +++++++---
 2 files changed, 18 insertions(+), 3 deletions(-)
 rename engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/{LibUtilsTest.java => ALibUtilsTest.java} (73%)

diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/LibUtils.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/LibUtils.java
index b28a7a9d19a..03835b6ca68 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/LibUtils.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/LibUtils.java
@@ -106,6 +106,17 @@ public static String getLibtorchPath() {
 
     private static void loadLibTorch(LibTorch libTorch) {
         Path libDir = libTorch.dir.toAbsolutePath();
+        if (Files.exists(libDir.resolve("libstdc++.so.6"))) {
+            String libstd = Utils.getEnvOrSystemProperty("LIBSTDCXX_LIBRARY_PATH");
+            if (libstd != null) {
+                try {
+                    logger.info("Loading libstdc++.so.6 from: {}", libstd);
+                    System.load(libstd);
+                } catch (UnsatisfiedLinkError e) {
+                    logger.warn("Failed Loading libstdc++.so.6 from: {}", libstd);
+                }
+            }
+        }
         boolean isCuda = libTorch.flavor.contains("cu");
         List<String> deferred =
                 Arrays.asList(
diff --git a/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/LibUtilsTest.java b/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/ALibUtilsTest.java
similarity index 73%
rename from engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/LibUtilsTest.java
rename to engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/ALibUtilsTest.java
index 617d2cfb809..f6cfda91106 100644
--- a/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/LibUtilsTest.java
+++ b/engines/pytorch/pytorch-engine/src/test/java/ai/djl/pytorch/integration/ALibUtilsTest.java
@@ -18,17 +18,21 @@
 import org.testng.annotations.BeforeClass;
 import org.testng.annotations.Test;
 
-public class LibUtilsTest {
+// Ensure this test run first
+public class ALibUtilsTest {
 
     @BeforeClass
     public void setup() {
-        System.setProperty(
-                "ai.djl.pytorch.native_helper", "ai.djl.pytorch.integration.LibUtilsTest");
+        System.setProperty("ai.djl.pytorch.native_helper", ALibUtilsTest.class.getName());
+        System.setProperty("STDCXX_LIBRARY_PATH", "/usr/lib/non-exists");
+        System.setProperty("PYTORCH_PRECXX11", "true");
     }
 
     @AfterClass
     public void teardown() {
         System.clearProperty("ai.djl.pytorch.native_helper");
+        System.clearProperty("LIBSTDCXX_LIBRARY_PATH");
+        System.clearProperty("PYTORCH_PRECXX11");
     }
 
     @Test