diff --git a/lib/torch-extension/arch.nix b/lib/torch-extension/arch.nix
index 20e72be7..0d5468c6 100644
--- a/lib/torch-extension/arch.nix
+++ b/lib/torch-extension/arch.nix
@@ -5,19 +5,25 @@
 
   lib,
   stdenv,
-  cudaPackages,
+
+  # Native build inputs
+  build2cmake,
   cmake,
   cmakeNvccThreadsHook,
-  ninja,
-  build2cmake,
   get-kernel-check,
   kernel-abi-check,
+  ninja,
   python3,
+  remove-bytecode-hook,
   rewrite-nix-paths-macho,
-  rocmPackages,
   writeScriptBin,
+
+  # Framework packages
+  cudaPackages,
+  rocmPackages,
   xpuPackages,
 
+  # Build inputs
   apple-sdk_15,
   clr,
   oneapi-torch-dev,
@@ -102,6 +108,7 @@ stdenv.mkDerivation (prevAttrs: {
     cmake
     ninja
     build2cmake
+    remove-bytecode-hook
   ]
   ++ lib.optionals doGetKernelCheck [
     get-kernel-check
diff --git a/overlay.nix b/overlay.nix
index 8907f6e7..97976395 100644
--- a/overlay.nix
+++ b/overlay.nix
@@ -11,6 +11,8 @@ final: prev: {
 
   rewrite-nix-paths-macho = prev.callPackage ./pkgs/rewrite-nix-paths-macho { };
 
+  remove-bytecode-hook = prev.callPackage ./pkgs/remove-bytecode-hook { };
+
   stdenvGlibc_2_27 = prev.callPackage ./pkgs/stdenv-glibc-2_27 { };
 
   # Python packages
diff --git a/pkgs/remove-bytecode-hook/default.nix b/pkgs/remove-bytecode-hook/default.nix
new file mode 100644
index 00000000..12c7f577
--- /dev/null
+++ b/pkgs/remove-bytecode-hook/default.nix
@@ -0,0 +1,5 @@
+{ makeSetupHook, python3 }:
+
+makeSetupHook {
+  name = "remove-bytecode-hook";
+} ./remove-bytecode-hook.sh
diff --git a/pkgs/remove-bytecode-hook/remove-bytecode-hook.sh b/pkgs/remove-bytecode-hook/remove-bytecode-hook.sh
new file mode 100755
index 00000000..47b6a180
--- /dev/null
+++ b/pkgs/remove-bytecode-hook/remove-bytecode-hook.sh
@@ -0,0 +1,12 @@
+#!/bin/sh
+
+echo "Sourcing remove-bytecode-hook.sh"
+
+removeBytecodeHook() {
+  echo "Removing Python bytecode"
+  find $out -type d -name '__pycache__' -exec rm -rf {} +
+}
+
+if [ -z "${dontRemoveBytecode-}" ]; then
+  appendToVar preDistPhases removeBytecodeHook
+fi