about summary refs log tree commit diff
path: root/pkgs/development/python-modules/torchinfo/default.nix
diff options
context:
space:
mode:
Diffstat (limited to 'pkgs/development/python-modules/torchinfo/default.nix')
-rw-r--r--pkgs/development/python-modules/torchinfo/default.nix30
1 files changed, 14 insertions, 16 deletions
diff --git a/pkgs/development/python-modules/torchinfo/default.nix b/pkgs/development/python-modules/torchinfo/default.nix
index 59b7e99c9dfc..a10dad4b10c0 100644
--- a/pkgs/development/python-modules/torchinfo/default.nix
+++ b/pkgs/development/python-modules/torchinfo/default.nix
@@ -1,12 +1,13 @@
-{ lib
-, buildPythonPackage
-, fetchFromGitHub
-, fetchpatch
-, pythonOlder
-, torch
-, torchvision
-, pytestCheckHook
-, transformers
+{
+  lib,
+  buildPythonPackage,
+  fetchFromGitHub,
+  fetchpatch,
+  pythonOlder,
+  torch,
+  torchvision,
+  pytestCheckHook,
+  transformers,
 }:
 
 buildPythonPackage rec {
@@ -24,12 +25,11 @@ buildPythonPackage rec {
   };
 
   patches = [
-    (fetchpatch {  # Add support for Python 3.11 and pytorch 2.1
+    (fetchpatch {
+      # Add support for Python 3.11 and pytorch 2.1
       url = "https://github.com/TylerYep/torchinfo/commit/c74784c71c84e62bcf56664653b7f28d72a2ee0d.patch";
       hash = "sha256-xSSqs0tuFpdMXUsoVv4sZLCeVnkK6pDDhX/Eobvn5mw=";
-      includes = [
-        "torchinfo/model_statistics.py"
-      ];
+      includes = [ "torchinfo/model_statistics.py" ];
     })
   ];
 
@@ -62,9 +62,7 @@ buildPythonPackage rec {
     "tests/torchinfo_xl_test.py"
   ];
 
-  pythonImportsCheck = [
-    "torchinfo"
-  ];
+  pythonImportsCheck = [ "torchinfo" ];
 
   meta = with lib; {
     description = "API to visualize pytorch models";