diff --git a/pkgs/development/libraries/science/math/tensorrt/extension.nix b/pkgs/development/libraries/science/math/tensorrt/extension.nix
index b4f47a8969c1..507dc0cfd505 100644
--- a/pkgs/development/libraries/science/math/tensorrt/extension.nix
+++ b/pkgs/development/libraries/science/math/tensorrt/extension.nix
@@ -27,7 +27,9 @@ final: prev: let
     defaultBuild = { "tensorrt" = if allBuilds ? ${computeName tensorRTDefaultVersion}
       then allBuilds.${computeName tensorRTDefaultVersion}
       else throw "tensorrt-${tensorRTDefaultVersion} does not support your cuda version ${cudaVersion}"; };
-  in allBuilds // defaultBuild;
+  in {
+    inherit buildTensorRTPackage;
+  } // allBuilds // defaultBuild;
 
   tensorRTVersions = {
     "8.4.0" = [