1{
2 lib,
3 buildPythonPackage,
4 fetchFromGitHub,
5 setuptools,
6 colorama,
7 matplotlib,
8 numpy,
9 pynvml,
10 torch,
11 torchprofile,
12}:
13
14buildPythonPackage {
15 pname = "pytorch-bench";
16 version = "unstable-2024-07-18";
17 pyproject = true;
18
19 src = fetchFromGitHub {
20 owner = "MaximeGloesener";
21 repo = "torch-benchmark";
22 rev = "405a3fc2d147b43b4c1f7edb7aca0a60ba343ac5";
23 hash = "sha256-KU3dAf97A6lkMNTKRay23BMFQfn1ReAFNaJ0kG2RfnA=";
24 };
25
26 build-system = [
27 setuptools
28 ];
29
30 dependencies = [
31 colorama
32 matplotlib
33 numpy
34 pynvml
35 torch
36 torchprofile
37 ];
38
39 pythonImportsCheck = [
40 "pytorch_bench"
41 ];
42
43 meta = {
44 description = "Benchmarking tool for torch";
45 homepage = "https://github.com/MaximeGloesener/torch-benchmark";
46 license = lib.licenses.mit;
47 maintainers = with lib.maintainers; [ ];
48 };
49}