diff --git a/pkgs/development/python-modules/albumentations/default.nix b/pkgs/development/python-modules/albumentations/default.nix
index 0c2949bbc068..6b6d4d08f715 100644
--- a/pkgs/development/python-modules/albumentations/default.nix
+++ b/pkgs/development/python-modules/albumentations/default.nix
@@ -12,6 +12,8 @@
 , pytestCheckHook
 , pythonOlder
 , pythonRelaxDepsHook
+, torch
+, torchvision
 }:
 
 buildPythonPackage rec {
@@ -50,6 +52,8 @@ buildPythonPackage rec {
   nativeCheckInputs = [
     deepdiff
     pytestCheckHook
+    torch
+    torchvision
   ];
 
   disabledTests = [