diff --git a/README.md b/README.md
new file mode 100644
--- /dev/null
+++ b/README.md
@@ -0,0 +1,134 @@
+# libtorch-ffi
+
+This package provides FFI bindings to PyTorch's libtorch C++ library.
+
+## Setup
+
+The package automatically downloads and configures libtorch during the build process. You can customize the setup using environment variables.
+
+### Environment Variables
+
+#### `LIBTORCH_VERSION`
+- **Default**: `2.5.0`
+- **Description**: Specifies the version of libtorch to download and use
+- **Example**: `export LIBTORCH_VERSION=2.5.0`
+
+#### `LIBTORCH_HOME`
+- **Default**: XDG cache directory (`~/.cache/libtorch` on Linux/macOS)
+- **Description**: Base directory where libtorch will be downloaded and stored
+- **Example**: `export LIBTORCH_HOME=/opt/libtorch`
+
+#### `LIBTORCH_CUDA_VERSION`
+- **Default**: `cpu`
+- **Description**: CUDA version for GPU support
+- **Options**: 
+  - `cpu` - CPU-only version (default)
+  - `cu117` - CUDA 11.7
+  - `cu118` - CUDA 11.8
+  - `cu121` - CUDA 12.1
+  - Any other CUDA version string supported by PyTorch
+- **Example**: `export LIBTORCH_CUDA_VERSION=cu118`
+
+#### `LIBTORCH_SKIP_DOWNLOAD`
+- **Default**: Not set
+- **Description**: When set (to any value), skips the automatic download of libtorch
+- **Use case**: When you have libtorch already installed system-wide
+- **Example**: `export LIBTORCH_SKIP_DOWNLOAD=1`
+
+### Directory Structure
+
+The downloaded libtorch is stored in a platform-specific directory structure:
+```
+$LIBTORCH_HOME/
+└── <version>/
+    └── <platform>/
+        └── <cuda-flavor>/
+            ├── lib/
+            ├── include/
+            └── .ok
+```
+
+Where:
+- `<version>` is the libtorch version (e.g., `2.5.0`)
+- `<platform>` is one of:
+  - `macos-arm64` - macOS on Apple Silicon
+  - `macos-x86_64` - macOS on Intel
+  - `linux-x86_64` - Linux on x86_64
+- `<cuda-flavor>` is the CUDA version (e.g., `cpu`, `cu118`)
+
+### Build Process
+
+1. **Pre-configuration**: The package checks if it's running in a Nix sandbox. If not, it proceeds with the download process.
+
+2. **Download**: If libtorch is not found in the cache directory, it will be automatically downloaded from PyTorch's official servers.
+
+3. **Configuration**: The build system automatically:
+   - Adds the libtorch library directory to the library search path
+   - Adds the include directories for C++ headers
+   - Sets up proper runtime library paths (rpath) for dynamic linking
+   - On macOS, adds the `-ld_classic` flag for compatibility
+
+### Platform-Specific Notes
+
+#### macOS
+- Uses rpath for dynamic library loading
+- Automatically adds `-ld_classic` flag for linker compatibility
+- Supports both Apple Silicon (arm64) and Intel (x86_64) architectures
+- **Since libtorch-ffi's rpath is propagated, it doesn't matter whether hasktorch is a static link or a shared link**
+
+#### Linux
+- Uses rpath for dynamic library loading
+- Supports x86_64 architecture
+- Multiple CUDA versions available for GPU support
+- **Since libtorch-ffi's rpath is not propagated, hasktorch must be a shared link**
+
+### Linking Configuration
+
+Due to rpath propagation differences between platforms, Linux requires shared linking. Add the following configuration:
+
+#### For Cabal (cabal.project)
+```
+shared: True
+executable-dynamic: True
+```
+
+#### For Stack (stack.yaml)
+```yaml
+configure-options:
+  $targets:
+    - --enable-executable-dynamic
+    - --enable-shared
+```
+
+### Nix Support
+
+The package detects when it's being built in a Nix sandbox and skips the automatic download. In this case, libtorch should be provided through Nix derivation inputs.
+
+### Troubleshooting
+
+1. **Download failures**: Check your internet connection and ensure the PyTorch download servers are accessible.
+
+2. **Missing libraries**: The `.ok` marker file indicates a successful download. If this file is missing but the directory exists, delete the directory and let the setup download again.
+
+3. **CUDA version mismatch**: Ensure your system CUDA version matches the `LIBTORCH_CUDA_VERSION` you've specified.
+
+4. **Custom libtorch installation**: Set `LIBTORCH_SKIP_DOWNLOAD=1` and ensure your system's libtorch is properly configured in your build environment.
+
+### Example Usage
+
+```bash
+# Use CPU-only version
+cabal build libtorch-ffi
+
+# Use CUDA 11.8 version
+export LIBTORCH_CUDA_VERSION=cu118
+cabal build libtorch-ffi
+
+# Use a specific version
+export LIBTORCH_VERSION=2.4.0
+cabal build libtorch-ffi
+
+# Use existing system libtorch
+export LIBTORCH_SKIP_DOWNLOAD=1
+cabal build libtorch-ffi
+```
diff --git a/Setup.hs b/Setup.hs
new file mode 100644
--- /dev/null
+++ b/Setup.hs
@@ -0,0 +1,198 @@
+import Data.List (isPrefixOf)
+import Distribution.Simple
+import Distribution.Simple.Program
+import Distribution.Simple.Setup
+import Distribution.Simple.LocalBuildInfo
+import Distribution.Types.LocalBuildInfo
+import Distribution.Types.GenericPackageDescription
+import Distribution.Types.HookedBuildInfo
+import Distribution.PackageDescription
+import Distribution.System
+import System.Directory
+import System.FilePath
+import System.Process
+import System.IO.Temp
+import Control.Monad
+import Network.HTTP.Simple
+import qualified Data.ByteString.Lazy as LBS
+import System.Environment (lookupEnv)
+import Data.Maybe (fromMaybe)
+import GHC.IO.Exception
+import Control.Exception
+
+main :: IO ()
+main = defaultMainWithHooks $ simpleUserHooks
+  { preConf = \_ _ -> do
+      pure emptyHookedBuildInfo
+  , confHook = \(gpd, hbi) flags -> do
+      mlibtorchDir <- ensureLibtorch
+      case mlibtorchDir of
+        Nothing -> do
+          putStrLn "libtorch not found, skipping configuration."
+          lbi <- confHook simpleUserHooks (gpd, hbi) flags
+          -- For macOS, add the -ld_classic flag to the linker
+          case buildOS of
+            OSX -> return $ lbi { withPrograms = addLdClassicFlag (withPrograms lbi) }
+            _ -> return $ lbi
+        Just libtorchDir -> do
+          libtorchDir <- getLocalUserLibtorchDir
+          let libDir     = libtorchDir </> "lib"
+              includeDir = libtorchDir </> "include"
+
+          let updatedFlags = flags
+                { configExtraLibDirs      = libDir : configExtraLibDirs flags
+                , configExtraIncludeDirs  =
+                    includeDir
+                    : (includeDir </> "torch" </> "csrc" </> "api" </> "include")
+                    : configExtraIncludeDirs flags
+                }
+          -- Call the default configuration hook with updated flags
+          lbi <- confHook simpleUserHooks (gpd, hbi) updatedFlags
+          -- For macOS, add the -ld_classic flag to the linker
+          case buildOS of
+            OSX -> return $ lbi { withPrograms = addRPath libDir $ addLdClassicFlag (withPrograms lbi) }
+            Linux -> return $ lbi { withPrograms = addRPath libDir (withPrograms lbi) }
+            _ -> return $ lbi
+  }
+
+getLibtorchVersion :: IO String
+getLibtorchVersion = do
+  mVersion <- lookupEnv "LIBTORCH_VERSION"
+  case mVersion of
+    Nothing -> return "2.5.0"
+    Just other -> return other
+
+getLocalUserLibtorchDir :: IO FilePath
+getLocalUserLibtorchDir = do
+  mHome <- lookupEnv "LIBTORCH_HOME"
+  libtorchVersion <- getLibtorchVersion
+  base <- case mHome of
+    Just h  -> pure h
+    Nothing -> do
+      -- XDG cache (Linux/macOS). Falls back to ~/.cache
+      cache <- getXdgDirectory XdgCache "libtorch"
+      pure cache
+  flavor <- getCudaFlavor
+  pure $ base </> libtorchVersion </> platformTag </> flavor
+
+platformTag :: FilePath
+platformTag =
+  case (buildOS, buildArch) of
+    (OSX,    AArch64) -> "macos-arm64"
+    (OSX,    X86_64)  -> "macos-x86_64"
+    (Linux,  X86_64)  -> "linux-x86_64"
+    -- add more as needed
+    _ -> error $ "Unsupported platform: " <> show (buildOS, buildArch)
+
+getCudaFlavor :: IO String
+getCudaFlavor = do
+  fromMaybe "cpu" <$> lookupEnv "LIBTORCH_CUDA_VERSION"  -- "cpu" | "cu117" | "cu118" | "cu121"
+
+ensureLibtorch :: IO (Maybe FilePath)
+ensureLibtorch = do
+  isSandbox <- isNixSandbox
+  if isSandbox
+    then return Nothing
+    else downloadLibtorch
+
+isNixSandbox :: IO Bool
+isNixSandbox = do
+  nix <- lookupEnv "NIX_BUILD_TOP"
+  case nix of
+    Just path -> do
+      let isNixPath = any (`isPrefixOf` path) ["/build", "/private/tmp/nix-build"]
+      if isNixPath
+        then do
+          putStrLn "Nix sandbox detected; skipping libtorch download."
+          return True
+        else do
+          return False
+    Nothing -> return False
+
+downloadLibtorch :: IO (Maybe FilePath)
+downloadLibtorch = do
+  skip <- lookupEnv "LIBTORCH_SKIP_DOWNLOAD"
+  case skip of
+    Just _ -> do
+      putStrLn "LIBTORCH_SKIP_DOWNLOAD set; assuming libtorch exists globally."
+      return Nothing
+    Nothing -> do
+      dest <- getLocalUserLibtorchDir
+      let marker = dest </> ".ok"
+      exists <- doesFileExist marker
+      present <- doesDirectoryExist dest
+      if present && exists
+        then pure $ Just dest
+        else do
+          putStrLn $ "libtorch not found in local cache, installing to " <> dest
+          downloadAndExtractLibtorchTo dest
+          -- Create an idempotence marker that checks
+          -- if we've already downloaded torch.
+          -- Since we'll be moving everything this will
+          -- be the our main reference.
+          writeFile marker ""
+          pure $ Just dest
+
+downloadAndExtractLibtorchTo :: FilePath -> IO ()
+downloadAndExtractLibtorchTo dest = do
+  createDirectoryIfMissing True dest
+  (url, fileName) <- computeURL
+  putStrLn $ "Downloading libtorch from: " ++ url
+  withSystemTempDirectory "libtorch-download" $ \tmpDir -> do
+    let downloadPath = tmpDir </> fileName
+    request  <- parseRequest url
+    response <- httpLBS request
+    LBS.writeFile downloadPath (getResponseBody response)
+    putStrLn "Download complete. Extracting..."
+    callProcess "unzip" ["-q", downloadPath, "-d", tmpDir]
+    let unpacked = tmpDir </> "libtorch"
+    exists <- doesDirectoryExist unpacked
+    let src = if exists then unpacked else tmpDir
+    -- We want to move the directory since this operation is atomic.
+    -- If that doesn't work we fall back to copying.
+    (renameDirectory src dest) `catch` (\(_::IOException) -> copyTree src dest)
+    putStrLn "libtorch extracted successfully (global cache)."
+
+copyTree :: FilePath -> FilePath -> IO ()
+copyTree src dest = do
+  createDirectoryIfMissing True dest
+  entries <- listDirectory src
+  mapM_ (\e -> do
+          let s = src  </> e
+              d = dest </> e
+          isDir <- doesDirectoryExist s
+          if isDir then copyTree s d else copyFile s d
+        ) entries
+
+computeURL :: IO (String, String)
+computeURL = do
+  flavor <- getCudaFlavor
+  v <- getLibtorchVersion
+  pure $ case buildOS of
+    OSX -> case buildArch of
+      AArch64 -> ( "https://download.pytorch.org/libtorch/cpu/libtorch-macos-arm64-" ++ v ++ ".zip"
+                 , "libtorch-macos-arm64.zip" )
+      X86_64  -> ( "https://download.pytorch.org/libtorch/cpu/libtorch-macos-x86_64-" ++ v ++ ".zip"
+                 , "libtorch-macos-x86_64.zip" )
+      _       -> error "Unsupported macOS arch"
+    Linux -> case flavor of
+      "cpu"  -> ( "https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-" ++ v ++ "%2Bcpu.zip"
+                , "libtorch-linux.zip" )
+      cudaVersion -> ( "https://download.pytorch.org/libtorch/" ++ cudaVersion ++"/libtorch-cxx11-abi-shared-with-deps-" ++ v ++ "%2B" ++ cudaVersion ++ ".zip"
+                , "libtorch-linux-" ++ cudaVersion ++ ".zip" )
+    Windows -> error "Windows not supported by this setup"
+
+-- Add -ld_classic flag to GHC program arguments for macOS
+addLdClassicFlag :: ProgramDb -> ProgramDb
+addLdClassicFlag progDb = 
+  case lookupProgram ghcProgram progDb of
+    Just ghc ->
+      let ghc' = ghc { programOverrideArgs = ["-optl-ld_classic"] ++ programOverrideArgs ghc }
+      in updateProgram ghc' progDb
+    Nothing -> progDb
+
+addRPath :: FilePath -> ProgramDb -> ProgramDb
+addRPath libDir progDb =
+  userSpecifyArgs (programName ldProgram)
+  ["-Wl,-rpath," ++ libDir]
+  progDb
diff --git a/libtorch-ffi.cabal b/libtorch-ffi.cabal
--- a/libtorch-ffi.cabal
+++ b/libtorch-ffi.cabal
@@ -1,6 +1,6 @@
 cabal-version:       3.0
 name:                libtorch-ffi
-version:             2.0.1.5
+version:             2.0.1.6
 -- The prefix(2.0) of this version("2.0.0.0") is the same as libtorch's one.
 synopsis:            Haskell bindings for PyTorch
 description:         This package provides Haskell bindings to libtorch, the C++ library underlying PyTorch, specifically designed for the Hasktorch ecosystem.
@@ -10,7 +10,19 @@
 maintainer:          hasktorch@gmail.com
 copyright:           2018 Austin Huang
 category:            Codegen
-build-type:          Simple
+build-type:          Custom
+extra-source-files:  README.md
+
+custom-setup
+  setup-depends:
+    base >= 4.7 && < 5,
+    Cabal >= 3.0 && < 4,
+    directory >= 1.3 && < 2,
+    filepath >= 1.4 && < 2,
+    process >= 1.6 && < 2,
+    temporary >= 1.3 && < 2,
+    http-conduit >= 2.3 && < 3,
+    bytestring >= 0.10 && < 1
 
 Flag cuda
  Description: A flag to link libtorch_cuda.