1 files changed, 33 insertions, 19 deletions

diff --git a/pkgs/development/python-modules/torch-tb-profiler/default.nix b/pkgs/development/python-modules/torch-tb-profiler/default.nix
index 108acfb0d355..766baff5d3cc 100644
--- a/pkgs/development/python-modules/torch-tb-profiler/default.nix
+++ b/pkgs/development/python-modules/torch-tb-profiler/default.nix
@@ -1,33 +1,43 @@
-{ buildPythonPackage
-, fetchFromGitHub
-, lib
-, pandas
-, pytestCheckHook
-, torch
-, tensorboard
-, torchvision
+{
+  lib,
+  buildPythonPackage,
+  fetchFromGitHub,
+  pandas,
+  pytestCheckHook,
+  setuptools,
+  tensorboard,
+  torch,
+  torchvision,
 }:
-
 let
-  version = "0.3.1";
+  version = "0.4.0";
   repo = fetchFromGitHub {
     owner = "pytorch";
     repo = "kineto";
-    rev = "v${version}";
-    hash = "sha256-Yg001XzOPDmz9wEP2b7Ggz/uU6x5PFzaaBeUBwWKFS0=";
+    rev = "refs/tags/v${version}";
+    hash = "sha256-nAtqGCv8q3Tati3NOGWWLb+gXdvO3qmECeC1WG2Mt3M=";
   };
 in
-buildPythonPackage rec {
+buildPythonPackage {
   pname = "torch_tb_profiler";
   inherit version;
-  format = "setuptools";
+  pyproject = true;
 
   # See https://discourse.nixos.org/t/extracting-sub-directory-from-fetchgit-or-fetchurl-or-any-derivation/8830.
   src = "${repo}/tb_plugin";
 
-  propagatedBuildInputs = [ pandas tensorboard ];
+  build-system = [ setuptools ];
+
+  dependencies = [
+    pandas
+    tensorboard
+  ];
 
-  nativeCheckInputs = [ pytestCheckHook torch torchvision ];
+  nativeCheckInputs = [
+    pytestCheckHook
+    torch
+    torchvision
+  ];
 
   disabledTests = [
     # Tests that attempt to access the filesystem in naughty ways.
@@ -35,14 +45,18 @@ buildPythonPackage rec {
     "test_tensorboard_end2end"
     "test_tensorboard_with_path_prefix"
     "test_tensorboard_with_symlinks"
+    "test_autograd_api"
+    "test_profiler_api_with_record_shapes_memory_stack"
+    "test_profiler_api_without_record_shapes_memory_stack"
+    "test_profiler_api_without_step"
   ];
 
   pythonImportsCheck = [ "torch_tb_profiler" ];
 
-  meta = with lib; {
+  meta = {
     description = "PyTorch Profiler TensorBoard Plugin";
     homepage = "https://github.com/pytorch/kineto";
-    license = licenses.bsd3;
-    maintainers = with maintainers; [ samuela ];
+    license = lib.licenses.bsd3;
+    maintainers = with lib.maintainers; [ samuela ];
   };
 }