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

   1 {
   2   lib,
   3   buildPythonPackage,
   4   fetchPypi,
   5   setuptools,
   6   torch,
   7 }:
   8
   9 buildPythonPackage rec {
  10   pname = "torchsummary";
  11   version = "1.5.1";
  12   pyproject = true;
  13
  14   src = fetchPypi {
  15     inherit pname version;
  16     hash = "sha256-mBv2ieIuDPf5XHRgAvIKJK0mqmudhhE0oUvGzpIjBZA=";
  17   };
  18
  19   build-system = [ setuptools ];
  20
  21   dependencies = [ torch ];
  22
  23   # no tests in pypi tarball
  24   doCheck = false;
  25
  26   pythonImportsCheck = [ "torchsummary" ];
  27
  28   meta = {
  29     description = "Model summary in PyTorch similar to `model.summary()` in Keras";
  30     homepage = "https://github.com/sksq96/pytorch-summary";
  31     license = lib.licenses.mit;
  32     maintainers = with lib.maintainers; [ tomasajt ];
  33   };
  34 }