pyrox.dev
1{
2 lib,
3 buildPythonPackage,
4 fetchFromGitHub,
5 setuptools,
6 numpy,
7 torch,
8 torchvision,
9}:
10
11buildPythonPackage rec {
12 pname = "torchprofile";
13 version = "0.0.4";
14 pyproject = true;
15
16 src = fetchFromGitHub {
17 owner = "zhijian-liu";
18 repo = "torchprofile";
19 tag = "v${version}";
20 hash = "sha256-ynRrGHzroyv8T8fggJAag7u6XBOx+uN49HSIe46Bcek=";
21 };
22
23 build-system = [
24 setuptools
25 ];
26
27 pythonRelaxDeps = [
28 "torchvision"
29 ];
30
31 dependencies = [
32 numpy
33 torch
34 torchvision
35 ];
36
37 pythonImportsCheck = [
38 "torchprofile"
39 ];
40
41 meta = {
42 changelog = "https://github.com/zhijian-liu/torchprofile/releases/tag/${src.tag}";
43 description = "General and accurate MACs / FLOPs profiler for PyTorch models";
44 homepage = "https://github.com/zhijian-liu/torchprofile";
45 license = lib.licenses.mit;
46 maintainers = with lib.maintainers; [ ];
47 };
48}