From d2f7acf30c6dc1fe3a742f4035b11fe57514952b Mon Sep 17 00:00:00 2001
From: Anders Kaseorg <andersk@mit.edu>
Date: Fri, 4 May 2018 18:00:50 -0400
Subject: [PATCH] pytorch: 0.3.1 -> 0.4.0

Signed-off-by: Anders Kaseorg <andersk@mit.edu>
---
 .../python-modules/pytorch/default.nix        | 35 +++++--------------
 pkgs/top-level/python-packages.nix            | 14 +++++++-
 2 files changed, 21 insertions(+), 28 deletions(-)

diff --git a/pkgs/development/python-modules/pytorch/default.nix b/pkgs/development/python-modules/pytorch/default.nix
index c8376196d3ac78..b0ac4e07093562 100644
--- a/pkgs/development/python-modules/pytorch/default.nix
+++ b/pkgs/development/python-modules/pytorch/default.nix
@@ -1,7 +1,7 @@
-{ buildPythonPackage,
+{ buildPythonPackage, pythonOlder,
   cudaSupport ? false, cudatoolkit ? null, cudnn ? null,
-  fetchFromGitHub, fetchpatch, lib, numpy, pyyaml, cffi, cmake,
-  git, stdenv, linkFarm, symlinkJoin,
+  fetchFromGitHub, fetchpatch, lib, numpy, pyyaml, cffi, typing, cmake,
+  stdenv, linkFarm, symlinkJoin,
   utillinux, which }:
 
 assert cudnn == null || cudatoolkit != null;
@@ -25,7 +25,7 @@ let
     "LD_LIBRARY_PATH=${cudaStub}\${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH} ";
 
 in buildPythonPackage rec {
-  version = "0.3.1";
+  version = "0.4.0";
   pname = "pytorch";
   name = "${pname}-${version}";
 
@@ -34,36 +34,17 @@ in buildPythonPackage rec {
     repo   = "pytorch";
     rev    = "v${version}";
     fetchSubmodules = true;
-    sha256 = "1k8fr97v5pf7rni5cr2pi21ixc3pdj3h3lkz28njbjbgkndh7mr3";
+    sha256 = "12d5vqqaprk0igmih7fwa65ldmaawgijxl58h6dnw660wysc132j";
   };
 
-  patches = [
-    (fetchpatch {
-      # make sure stdatomic.h is included when checking for ATOMIC_INT_LOCK_FREE
-      # Fixes this test failure:
-      # RuntimeError: refcounted file mapping not supported on your system at /tmp/nix-build-python3.6-pytorch-0.3.0.drv-0/source/torch/lib/TH/THAllocator.c:525
-      url = "https://github.com/pytorch/pytorch/commit/502aaf39cf4a878f9e4f849e5f409573aa598aa9.patch";
-      stripLen = 3;
-      extraPrefix = "torch/lib/";
-      sha256 = "1miz4lhy3razjwcmhxqa4xmlcmhm65lqyin1czqczj8g16d3f62f";
-    })
-  ];
-
-  postPatch = ''
-    substituteInPlace test/run_test.sh --replace \
-      "INIT_METHOD='file://'\$TEMP_DIR'/shared_init_file' \$PYCMD ./test_distributed.py" \
-      "echo Skipped for Nix package"
-  '';
-
   preConfigure = lib.optionalString cudaSupport ''
-    export CC=${cudatoolkit.cc}/bin/gcc
+    export CC=${cudatoolkit.cc}/bin/gcc CXX=${cudatoolkit.cc}/bin/g++
   '' + lib.optionalString (cudaSupport && cudnn != null) ''
     export CUDNN_INCLUDE_DIR=${cudnn}/include
   '';
 
   buildInputs = [
      cmake
-     git
      numpy.blas
      utillinux
      which
@@ -73,10 +54,10 @@ in buildPythonPackage rec {
     cffi
     numpy
     pyyaml
-  ];
+  ] ++ lib.optional (pythonOlder "3.5") typing;
 
   checkPhase = ''
-    ${cudaStubEnv}${stdenv.shell} test/run_test.sh
+    ${cudaStubEnv}python test/run_test.py --exclude distributed
   '';
 
   meta = {
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index 7868f7d7ff06eb..95b7f8a5718cdd 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -5623,8 +5623,20 @@ in {
     };
   };
 
-  pytorch = callPackage ../development/python-modules/pytorch {
+  pytorch = let
+    # Fails with CUDA 9.1 and GCC 6.4:
+    # https://github.com/pytorch/pytorch/issues/5831
+    # https://devtalk.nvidia.com/default/topic/1028112
+    # We should be able to remove this when CUDA 9.2 is released.
+    cudatoolkit9 = pkgs.cudatoolkit9.override {
+      gcc6 = pkgs.gcc5;
+    };
+  in callPackage ../development/python-modules/pytorch {
     cudaSupport = pkgs.config.cudaSupport or false;
+    cudatoolkit = cudatoolkit9;
+    cudnn = pkgs.cudnn_cudatoolkit9.override {
+      inherit cudatoolkit9;
+    };
   };
 
   pytorchWithCuda = self.pytorch.override {