pyrox.dev
1{ lib
2, buildPythonPackage
3, fetchFromGitHub
4, pythonOlder
5, pytestCheckHook
6, torch
7}:
8
9buildPythonPackage 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}