pkgs/development/python-modules/torchinfo/default.nix

   1 { lib
   2 , buildPythonPackage
   3 , fetchFromGitHub
   4 , pytestCheckHook
   5 , pythonOlder
   6 , torch
   7 , torchvision
   8 }:
   9
  10 buildPythonPackage rec {
  11   pname = "torchinfo";
  12   version = "1.64";
  13   format = "setuptools";
  14
  15   disabled = pythonOlder "3.7";
  16
  17   src = fetchFromGitHub {
  18     owner = "TylerYep";
  19     repo = pname;
  20     rev = "refs/tags/v${version}";
  21     hash = "sha256-gcl8RxCD017FP4LtB60WVtOh7jg2Otv/vNd9hKneEAU=";
  22   };
  23
  24   propagatedBuildInputs = [
  25     torch
  26     torchvision
  27   ];
  28
  29   checkInputs = [
  30     pytestCheckHook
  31   ];
  32
  33   disabledTests = [
  34     # Skip as it downloads pretrained weights (require network access)
  35     "test_eval_order_doesnt_matter"
  36     # AssertionError in output
  37     "test_google"
  38   ];
  39
  40   pythonImportsCheck = [
  41     "torchvision"
  42   ];
  43
  44   meta = with lib; {
  45     description = "API to visualize pytorch models";
  46     homepage = "https://github.com/TylerYep/torchinfo";
  47     license = licenses.mit;
  48     maintainers = with maintainers; [ petterstorvik ];
  49   };
  50 }