From f803fa80cc1d32370498d561c150abb88bb41d38 Mon Sep 17 00:00:00 2001
From: Pol Dellaiera <pol.dellaiera@protonmail.com>
Date: Fri, 1 Nov 2024 11:08:19 +0100
Subject: [PATCH] python312Packages.pytorch-bench: init at unstable-2024-07-18

---
 .../python-modules/pytorch-bench/default.nix  | 49 +++++++++++++++++++
 pkgs/top-level/python-packages.nix            |  2 +
 2 files changed, 51 insertions(+)
 create mode 100644 pkgs/development/python-modules/pytorch-bench/default.nix

diff --git a/pkgs/development/python-modules/pytorch-bench/default.nix b/pkgs/development/python-modules/pytorch-bench/default.nix
new file mode 100644
index 0000000000000..e705028b1db8f
--- /dev/null
+++ b/pkgs/development/python-modules/pytorch-bench/default.nix
@@ -0,0 +1,49 @@
+{
+  lib,
+  buildPythonPackage,
+  fetchFromGitHub,
+  setuptools,
+  colorama,
+  matplotlib,
+  numpy,
+  pynvml,
+  torch,
+  torchprofile,
+}:
+
+buildPythonPackage {
+  pname = "pytorch-bench";
+  version = "unstable-2024-07-18";
+  pyproject = true;
+
+  src = fetchFromGitHub {
+    owner = "MaximeGloesener";
+    repo = "torch-benchmark";
+    rev = "405a3fc2d147b43b4c1f7edb7aca0a60ba343ac5";
+    hash = "sha256-KU3dAf97A6lkMNTKRay23BMFQfn1ReAFNaJ0kG2RfnA=";
+  };
+
+  build-system = [
+    setuptools
+  ];
+
+  dependencies = [
+    colorama
+    matplotlib
+    numpy
+    pynvml
+    torch
+    torchprofile
+  ];
+
+  pythonImportsCheck = [
+    "pytorch_bench"
+  ];
+
+  meta = {
+    description = "Benchmarking tool for torch";
+    homepage = "https://github.com/MaximeGloesener/torch-benchmark";
+    license = lib.licenses.mit;
+    maintainers = with lib.maintainers; [ drupol ];
+  };
+}
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index dec0803ece86a..b68012dcf33f6 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -13093,6 +13093,8 @@ self: super: with self; {
 
   pytools = callPackage ../development/python-modules/pytools { };
 
+  pytorch-bench = callPackage ../development/python-modules/pytorch-bench { };
+
   pytorch-lightning = callPackage ../development/python-modules/pytorch-lightning { };
 
   pytorch-metric-learning = callPackage ../development/python-modules/pytorch-metric-learning { };