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

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