tjh.dev
1{ stdenv
2, buildPythonPackage
3, fetchurl
4, isPy37
5, isPy38
6, python
7, nvidia_x11
8, addOpenGLRunpath
9, future
10, numpy
11, patchelf
12, pyyaml
13, requests
14}:
15
16let
17 pyVerNoDot = builtins.replaceStrings [ "." ] [ "" ] python.pythonVersion;
18 platform = if stdenv.isDarwin then "darwin" else "linux";
19 srcs = import ./binary-hashes.nix;
20 unsupported = throw "Unsupported system";
21in buildPythonPackage {
22 pname = "pytorch";
23 # Don't forget to update pytorch to the same version.
24 version = "1.6.0";
25
26 format = "wheel";
27
28 disabled = !(isPy37 || isPy38);
29
30 src = fetchurl srcs."${stdenv.system}-${pyVerNoDot}" or unsupported;
31
32 nativeBuildInputs = [
33 addOpenGLRunpath
34 patchelf
35 ];
36
37 propagatedBuildInputs = [
38 future
39 numpy
40 pyyaml
41 requests
42 ];
43
44 postInstall = ''
45 # ONNX conversion
46 rm -rf $out/bin
47 '';
48
49 postFixup = let
50 rpath = stdenv.lib.makeLibraryPath [ stdenv.cc.cc.lib nvidia_x11 ];
51 in ''
52 find $out/${python.sitePackages}/torch/lib -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
53 echo "setting rpath for $lib..."
54 patchelf --set-rpath "${rpath}:$out/${python.sitePackages}/torch/lib" "$lib"
55 addOpenGLRunpath "$lib"
56 done
57 '';
58
59 pythonImportsCheck = [ "torch" ];
60
61 meta = with stdenv.lib; {
62 description = "Open source, prototype-to-production deep learning platform";
63 homepage = "https://pytorch.org/";
64 license = licenses.unfree; # Includes CUDA and Intel MKL.
65 platforms = platforms.linux;
66 maintainers = with maintainers; [ danieldk ];
67 };
68}