1{
2 lib,
3 buildPythonPackage,
4 fetchFromGitHub,
5 distutils,
6 fsspec,
7 graphviz,
8 torch,
9 unittestCheckHook,
10}:
11
12buildPythonPackage rec {
13 pname = "pytorchviz";
14 version = "0.0.2-unstable-2024-12-30";
15 pyproject = true;
16
17 src = fetchFromGitHub {
18 owner = "szagoruyko";
19 repo = "pytorchviz";
20 # No tags in the upstream GitHub repo
21 rev = "5cf04c13e601366f6b9cf5939b5af5144d55b887";
22 hash = "sha256-La1X8Y64n/vNGDUEsw1iZ5Mb6/w3WayeWxa62QxLyHA=";
23 };
24
25 dependencies = [
26 graphviz
27 torch
28 ];
29
30 nativeCheckInputs = [
31 unittestCheckHook
32 distutils
33 fsspec
34 ];
35
36 unittestFlagsArray = [ "test/" ];
37
38 pythonImportsCheck = [ "torchviz" ];
39
40 meta = {
41 description = "Small package to create visualizations of PyTorch execution graphs";
42 homepage = "https://github.com/szagoruyko/pytorchviz";
43 license = lib.licenses.mit;
44 maintainers = with lib.maintainers; [ jherland ];
45 };
46}