pyrox.dev
1{
2 lib,
3 buildPythonPackage,
4 fetchFromGitHub,
5 calmsize,
6 pandas,
7 torch,
8 pytestCheckHook,
9 setuptools,
10}:
11
12buildPythonPackage rec {
13 pname = "pytorch-memlab";
14 version = "0.3.0";
15 pyproject = true;
16
17 src = fetchFromGitHub {
18 owner = "Stonesjtu";
19 repo = "pytorch_memlab";
20 tag = version;
21 hash = "sha256-wNgbipvi3vYr9Ka9hA7I+C4y8Nf6AiZXUoXX+qKtZ+I=";
22 };
23
24 build-system = [ setuptools ];
25
26 nativeCheckInputs = [ pytestCheckHook ];
27
28 dependencies = [
29 calmsize
30 pandas
31 torch
32 ];
33
34 pythonImportsCheck = [ "pytorch_memlab" ];
35
36 # These tests require CUDA
37 disabledTestPaths = [
38 "test/test_courtesy.py"
39 "test/test_line_profiler.py"
40 ];
41
42 meta = {
43 description = "Simple and accurate CUDA memory management laboratory for pytorch";
44 homepage = "https://github.com/Stonesjtu/pytorch_memlab";
45 changelog = "https://github.com/Stonesjtu/pytorch_memlab/releases/tag/${version}";
46 license = lib.licenses.mit;
47 maintainers = with lib.maintainers; [ jherland ];
48 };
49}