1 files changed, 20 insertions, 3 deletions
diff --git a/infra/libkookie/nixpkgs/pkgs/development/python-modules/pytorch-metric-learning/default.nix b/infra/libkookie/nixpkgs/pkgs/development/python-modules/pytorch-metric-learning/default.nix
index daf6accb8935..64564a4dca55 100644
--- a/infra/libkookie/nixpkgs/pkgs/development/python-modules/pytorch-metric-learning/default.nix
+++ b/infra/libkookie/nixpkgs/pkgs/development/python-modules/pytorch-metric-learning/default.nix
@@ -4,6 +4,7 @@
 , isPy27
 , numpy
 , scikitlearn
+, pytestCheckHook
 , pytorch
 , torchvision
 , tqdm
@@ -11,15 +12,15 @@
 
 buildPythonPackage rec {
   pname   = "pytorch-metric-learning";
-  version = "0.9.81";
+  version = "0.9.94";
 
   disabled = isPy27;
 
   src = fetchFromGitHub {
     owner = "KevinMusgrave";
     repo = pname;
-    rev = "cb23328aba64f7f4658374cc2920ef5d56cda5c8";  # no version tag
-    sha256 = "0c2dyi4qi7clln43481xq66f6r4fadrz84jphjc5phz97bp33ds8";
+    rev = "v${version}";
+    sha256 = "1i2m651isa6xk3zj8dhzdbmd1bdzl51bh6rxifx6gg22hfa5dj9a";
   };
 
   propagatedBuildInputs = [
@@ -30,6 +31,22 @@ buildPythonPackage rec {
     tqdm
   ];
 
+  preCheck = ''
+    export HOME=$TMP
+    export TEST_DEVICE=cpu
+    export TEST_DTYPES=float32,float64  # half-precision tests fail on CPU
+  '';
+  # package only requires `unittest`, but use `pytest` to exclude tests
+  checkInputs = [ pytestCheckHook ];
+  disabledTests = [
+    # requires FAISS (not in Nixpkgs)
+    "test_accuracy_calculator_and_faiss"
+    # require network access:
+    "test_get_nearest_neighbors"
+    "test_tuplestoweights_sampler"
+    "test_untrained_indexer"
+  ];
+
   meta = {
     description = "Metric learning library for PyTorch";
     homepage = "https://github.com/KevinMusgrave/pytorch-metric-learning";