blob: 475e7627e4baa215fec436fd8ab748f2eaebffe1 (
plain) (
blame)
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
|
{ lib
, python
, buildPythonPackage
, autoPatchelfHook
, unzip
, cudaPackages
}:
let
pyVersion = "${lib.versions.major python.version}${lib.versions.minor python.version}";
in
buildPythonPackage rec {
pname = "tensorrt";
version = cudaPackages.tensorrt.version;
src = cudaPackages.tensorrt.src;
format = "wheel";
# We unpack the wheel ourselves because of the odd packaging.
dontUseWheelUnpack = true;
nativeBuildInputs = [
unzip
autoPatchelfHook
cudaPackages.autoAddOpenGLRunpathHook
];
preUnpack = ''
mkdir -p dist
tar --strip-components=2 -xf "$src" --directory=dist \
"TensorRT-${version}/python/tensorrt-${version}-cp${pyVersion}-none-linux_x86_64.whl"
'';
sourceRoot = ".";
buildInputs = [
cudaPackages.cudnn
cudaPackages.tensorrt
];
pythonImportsCheck = [
"tensorrt"
];
meta = with lib; {
description = "Python bindings for TensorRT, a high-performance deep learning interface";
homepage = "https://developer.nvidia.com/tensorrt";
license = licenses.unfree;
platforms = [ "x86_64-linux" ];
maintainers = with maintainers; [ aidalgol ];
};
}
|