1 files changed, 5 insertions, 7 deletions
diff --git a/infra/libkookie/nixpkgs/pkgs/development/python-modules/torchgpipe/default.nix b/infra/libkookie/nixpkgs/pkgs/development/python-modules/torchgpipe/default.nix
index 6e621ee8b688..dbfad336d0f3 100644
--- a/infra/libkookie/nixpkgs/pkgs/development/python-modules/torchgpipe/default.nix
+++ b/infra/libkookie/nixpkgs/pkgs/development/python-modules/torchgpipe/default.nix
@@ -2,7 +2,6 @@
 , buildPythonPackage
 , fetchFromGitHub
 , isPy27
-, pytest
 , pytestrunner
 , pytestCheckHook
 , pytorch
@@ -23,12 +22,11 @@ buildPythonPackage rec {
 
   propagatedBuildInputs = [ pytorch ];
 
-  checkInputs = [ pytest pytestrunner pytestCheckHook ];
-  disabledTests = [ "test_inplace_on_requires_grad" ];
-  # seems like a harmless failure:
-  ## AssertionError:
-  ## Pattern 'a leaf Variable that requires grad has been used in an in-place operation.'
-  ## does not match 'a leaf Variable that requires grad is being used in an in-place operation.'
+  checkInputs = [ pytestrunner pytestCheckHook ];
+  disabledTests = [
+    "test_inplace_on_requires_grad"
+    "test_input_requiring_grad"
+  ];
 
   meta = with lib; {
     description = "GPipe implemented in Pytorch and optimized for CUDA rather than TPU";