blob: 60f9b5ad88464be5abb805c8c8db17180a3bdb4a (
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
|
{ lib
, stdenv
, cmake
, libtorch-bin
, linkFarm
, symlinkJoin
, cudaSupport
, cudatoolkit
, cudnn
}:
let
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 = ./.;
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 = ''
LD_LIBRARY_PATH=${cudaStub}''${LD_LIBRARY_PATH:+:}$LD_LIBRARY_PATH \
./test
'';
}
|