diff --git a/pkgs/development/python-modules/pytorch/default.nix b/pkgs/development/python-modules/pytorch/default.nix new file mode 100644 index 00000000000..2088b2e0b10 --- /dev/null +++ b/pkgs/development/python-modules/pytorch/default.nix @@ -0,0 +1,43 @@ +{ buildPythonPackage, fetchFromGitHub, lib, numpy, pyyaml, cffi, cmake, + git, stdenv }: + +buildPythonPackage rec { + version = "0.1.12"; + pname = "pytorch"; + name = "${pname}-${version}"; + + src = fetchFromGitHub { + owner = "pytorch"; + repo = "pytorch"; + rev = "v${version}"; + sha256 = "0r8mf4xya76gz83y5z3hfxh0rydkydafhipl8g7d0bfrgw961jy9"; + }; + + checkPhase = '' + ${stdenv.shell} test/run_test.sh + ''; + + buildInputs = [ + cmake + git + numpy.blas + ]; + + propagatedBuildInputs = [ + cffi + numpy + pyyaml + ]; + + preConfigure = '' + export NO_CUDA=1 + ''; + + meta = { + description = "Tensors and Dynamic neural networks in Python with strong GPU acceleration."; + homepage = http://pytorch.org/; + license = lib.licenses.bsd3; + platforms = lib.platforms.linux; + maintainers = with lib.maintainers; [ teh ]; + }; +} diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix index d79ba6af0e1..2afe071aa66 100644 --- a/pkgs/top-level/python-packages.nix +++ b/pkgs/top-level/python-packages.nix @@ -10762,6 +10762,8 @@ in { }; }; + pytorch = callPackage ../development/python-modules/pytorch { }; + python_tvrage = buildPythonPackage (rec { version = "0.4.1"; name = "tvrage-${version}";