pyrox.dev
1{ lib
2, buildPythonPackage
3, fetchFromGitHub
4, isPy27
5, pytest-runner
6, pytestCheckHook
7, pytorch
8}:
9
10buildPythonPackage 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 = [ pytest-runner 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}