1 files changed, 39 insertions, 33 deletions
diff --git a/pkgs/development/python-modules/tensordict/default.nix b/pkgs/development/python-modules/tensordict/default.nix
index 6c64ca00c5ee5..bd35eff65707e 100644
--- a/pkgs/development/python-modules/tensordict/default.nix
+++ b/pkgs/development/python-modules/tensordict/default.nix
@@ -1,21 +1,23 @@
-{ lib
-, buildPythonPackage
-, pythonOlder
-, fetchFromGitHub
-, setuptools
-, torch
-, wheel
-, which
-, cloudpickle
-, numpy
-, h5py
-, pytestCheckHook
-, stdenv
+{
+  lib,
+  buildPythonPackage,
+  pythonOlder,
+  fetchFromGitHub,
+  setuptools,
+  torch,
+  wheel,
+  which,
+  cloudpickle,
+  numpy,
+  h5py,
+  pytestCheckHook,
+  stdenv,
+  pythonAtLeast,
 }:
 
 buildPythonPackage rec {
   pname = "tensordict";
-  version = "0.3.1";
+  version = "0.4.0";
   pyproject = true;
 
   disabled = pythonOlder "3.8";
@@ -24,25 +26,23 @@ buildPythonPackage rec {
     owner = "pytorch";
     repo = "tensordict";
     rev = "refs/tags/v${version}";
-    hash = "sha256-eCx1r7goqOdGX/0mSGCiLhdGQTh4Swa5aFiLSsL56p0=";
+    hash = "sha256-wKEzNaaazGEkoElzp93RIlq/r5uRUdM7UyDy/DygIEc=";
   };
 
-  nativeBuildInputs = [
+  build-system = [
     setuptools
     torch
     wheel
     which
   ];
 
-  propagatedBuildInputs = [
+  dependencies = [
     cloudpickle
     numpy
     torch
   ];
 
-  pythonImportsCheck = [
-    "tensordict"
-  ];
+  pythonImportsCheck = [ "tensordict" ];
 
   # We have to delete the source because otherwise it is used instead of the installed package.
   preCheck = ''
@@ -54,23 +54,29 @@ buildPythonPackage rec {
     pytestCheckHook
   ];
 
-  # RuntimeError: internal error
-  disabledTests = lib.optionals (stdenv.hostPlatform.system == "aarch64-linux") [
-    "test_add_scale_sequence"
-    "test_modules"
-    "test_setattr"
-  ];
+  disabledTests =
+    # Hangs forever
+    [ "test_copy_onto" ]
+    ++ lib.optionals (stdenv.hostPlatform.system == "aarch64-linux") [
+      # RuntimeError: internal error
+      "test_add_scale_sequence"
+      "test_modules"
+      "test_setattr"
+
+      # _queue.Empty errors in multiprocessing tests
+      "test_isend"
+    ];
 
   # ModuleNotFoundError: No module named 'torch._C._distributed_c10d'; 'torch._C' is not a package
-  disabledTestPaths = lib.optionals stdenv.isDarwin [
-    "test/test_distributed.py"
-  ];
+  disabledTestPaths = lib.optionals stdenv.isDarwin [ "test/test_distributed.py" ];
 
-  meta = with lib; {
-    description = "A pytorch dedicated tensor container";
+  meta = {
+    description = "Pytorch dedicated tensor container";
     changelog = "https://github.com/pytorch/tensordict/releases/tag/v${version}";
     homepage = "https://github.com/pytorch/tensordict";
-    license = licenses.mit;
-    maintainers = with maintainers; [ GaetanLepage ];
+    license = lib.licenses.mit;
+    maintainers = with lib.maintainers; [ GaetanLepage ];
+    # No python 3.12 support yet: https://github.com/pytorch/rl/issues/2035
+    broken = pythonAtLeast "3.12";
   };
 }