pyrox.dev
1{ buildPythonPackage, fetchFromGitHub, lib, numpy, pyyaml, cffi, cmake,
2 git, stdenv }:
3
4buildPythonPackage rec {
5 version = "0.2.0";
6 pname = "pytorch";
7 name = "${pname}-${version}";
8
9 src = fetchFromGitHub {
10 owner = "pytorch";
11 repo = "pytorch";
12 rev = "v${version}";
13 sha256 = "1s3f46ga1f4lfrcj3lpvvhgkdr1pi8i2hjd9xj9qiz3a9vh2sj4n";
14 };
15
16 checkPhase = ''
17 ${stdenv.shell} test/run_test.sh
18 '';
19
20 buildInputs = [
21 cmake
22 git
23 numpy.blas
24 ];
25
26 propagatedBuildInputs = [
27 cffi
28 numpy
29 pyyaml
30 ];
31
32 preConfigure = ''
33 export NO_CUDA=1
34 '';
35
36 meta = {
37 description = "Tensors and Dynamic neural networks in Python with strong GPU acceleration.";
38 homepage = http://pytorch.org/;
39 license = lib.licenses.bsd3;
40 platforms = lib.platforms.linux;
41 maintainers = with lib.maintainers; [ teh ];
42 };
43}