Skip to content

Commit

Permalink
[pytorch] Avoid loading .lib file from PYTORCH_LIBRARY_PATH (#2987)
Browse files Browse the repository at this point in the history
  • Loading branch information
frankfliu committed Apr 26, 2024
1 parent d0ad200 commit 3a37219
Showing 1 changed file with 4 additions and 1 deletion.
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,7 @@ public final class LibUtils {

private static final Pattern VERSION_PATTERN =
Pattern.compile("(\\d+\\.\\d+\\.\\d+(-[a-z]+)?)(-SNAPSHOT)?(-\\d+)?");
private static final Pattern LIB_PATTERN = Pattern.compile("(.*\\.(so(\\.\\d+)*|dll|dylib))");

private static LibTorch libTorch;

Expand Down Expand Up @@ -136,7 +137,9 @@ private static void loadLibTorch(LibTorch libTorch) {
paths.filter(
path -> {
String name = path.getFileName().toString();
if (!isCuda
if (!LIB_PATTERN.matcher(name).matches()) {
return false;
} else if (!isCuda
&& name.contains("nvrtc")
&& name.contains("cudart")
&& name.contains("nvTools")) {
Expand Down

0 comments on commit 3a37219

Please sign in to comment.