pkgs/development/python-modules/torchgpipe/default.nix

   1 { lib
   2 , buildPythonPackage
   3 , fetchFromGitHub
   4 , pythonOlder
   5 , pytestCheckHook
   6 , torch
   7 }:
   8
   9 buildPythonPackage rec {
  10   pname = "torchgpipe";
  11   version = "0.0.7";
  12
  13   disabled = pythonOlder "3.6";
  14
  15   src = fetchFromGitHub {
  16     owner = "kakaobrain";
  17     repo = pname;
  18     rev = "v${version}";
  19     sha256 = "0ki0njhmz1i3pkpr3y6h6ac7p5qh1kih06mknc2s18mfw34f2l55";
  20   };
  21
  22   postPatch = ''
  23     substituteInPlace setup.py \
  24       --replace "'pytest-runner'" ""
  25   '';
  26
  27   propagatedBuildInputs = [ torch ];
  28
  29   nativeCheckInputs = [ pytestCheckHook ];
  30   disabledTests = [
  31     "test_inplace_on_requires_grad"
  32     "test_input_requiring_grad"
  33   ];
  34
  35   meta = with lib; {
  36     description = "GPipe implemented in Pytorch and optimized for CUDA rather than TPU";
  37     homepage = "https://torchgpipe.readthedocs.io";
  38     license = licenses.asl20;
  39     maintainers = [ maintainers.bcdarwin ];
  40   };
  41 }