blob: 01ae0e7d2c9756fce6b5d946223516a08448cbba (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
|
{
lib,
stdenv,
cmake,
libtorch-bin,
linkFarm,
symlinkJoin,
cudaSupport,
cudaPackages ? { },
}:
let
inherit (cudaPackages) cudatoolkit cudnn;
cudatoolkit_joined = symlinkJoin {
name = "${cudatoolkit.name}-unsplit";
paths = [
cudatoolkit.out
cudatoolkit.lib
];
};
# We do not have access to /run/opengl-driver/lib in the sandbox,
# so use a stub instead.
cudaStub = linkFarm "cuda-stub" [
{
name = "libcuda.so.1";
path = "${cudatoolkit}/lib/stubs/libcuda.so";
}
];
in
stdenv.mkDerivation {
pname = "libtorch-test";
version = libtorch-bin.version;
src = lib.fileset.toSource {
root = ./.;
fileset = lib.fileset.unions [
./CMakeLists.txt
./test.cpp
];
};
nativeBuildInputs = [ cmake ];
buildInputs = [ libtorch-bin ] ++ lib.optionals cudaSupport [ cudnn ];
cmakeFlags = lib.optionals cudaSupport [ "-DCUDA_TOOLKIT_ROOT_DIR=${cudatoolkit_joined}" ];
doCheck = true;
installPhase = ''
touch $out
'';
checkPhase =
lib.optionalString cudaSupport ''
LD_LIBRARY_PATH=${cudaStub}''${LD_LIBRARY_PATH:+:}$LD_LIBRARY_PATH \
''
+ ''
./test
'';
}
|
