diff --git a/CHANGELOG.md b/CHANGELOG.md
new file mode 100644
--- /dev/null
+++ b/CHANGELOG.md
@@ -0,0 +1,8 @@
+# Changelog for libtorch-ffi
+
+## 2.0.2.0
+
+- Add support for Stackage lts-23.24 (GHC 9.8); the previous release used GHC 9.6.
+- Expose parameter group interface in `Torch.Internal.Managed.Optim` and `Torch.Internal.Unmanaged.Optim` (supports per-group Adam/AdamW options).
+- Implement non-blocking host-to-device transfer.
+- Setup: dynamic detection for CUDA package installation, download NVIDIA libs, drop `rpath-link` on macOS.
diff --git a/Setup.hs b/Setup.hs
--- a/Setup.hs
+++ b/Setup.hs
@@ -1,6 +1,9 @@
 {-# LANGUAGE CPP #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 
-import Data.List (isPrefixOf)
+import Data.Char (toLower)
+import Data.List (isPrefixOf, isInfixOf, nubBy)
 import Distribution.Simple
 import Distribution.Simple.Program
 import Distribution.Simple.Setup
@@ -17,8 +20,9 @@
 import Control.Monad
 import Network.HTTP.Simple
 import qualified Data.ByteString.Lazy as LBS
+import qualified Data.ByteString.Lazy.Char8 as LBSC
 import System.Environment (lookupEnv)
-import Data.Maybe (fromMaybe)
+import Data.Maybe (fromMaybe, listToMaybe)
 import GHC.IO.Exception
 import Control.Exception
 import Codec.Archive.Zip
@@ -39,7 +43,7 @@
       mlibtorchDir <- ensureLibtorch
       case mlibtorchDir of
         Nothing -> do
-          putStrLn "libtorch not found, skipping configuration."
+          putStrLn "libtorch not found or handled by Nix, skipping configuration."
           lbi <- confHook simpleUserHooks (gpd, hbi) flags
           -- For macOS, add the -ld_classic flag to the linker
           case buildOS of
@@ -59,7 +63,8 @@
                 }
           -- Call the default configuration hook with updated flags
           lbi <- confHook simpleUserHooks (gpd, hbi) updatedFlags
-          -- For macOS, add the -ld_classic flag to the linker
+          
+          -- Add RPath so the binary finds the libs at runtime
           case buildOS of
             OSX -> return $ lbi { withPrograms = addRPath libDir $ addLdClassicFlag (withPrograms lbi) }
             Linux -> return $ lbi { withPrograms = addRPath libDir (withPrograms lbi) }
@@ -97,7 +102,7 @@
 
 getCudaFlavor :: IO String
 getCudaFlavor = do
-  fromMaybe "cpu" <$> lookupEnv "LIBTORCH_CUDA_VERSION"  -- "cpu" | "cu126" | "cu128" | "cu130"
+  fromMaybe "cpu" <$> lookupEnv "LIBTORCH_CUDA_VERSION"  -- "cpu" | "cu121" | "cu118"
 
 ensureLibtorch :: IO (Maybe FilePath)
 ensureLibtorch = do
@@ -137,10 +142,15 @@
         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.
+
+          -- Download NVIDIA libs if on Linux + CUDA
+          flavor <- getCudaFlavor
+          when (buildOS == Linux && "cu" `isPrefixOf` flavor) $ do
+             -- Dynamically detect which NVIDIA library versions are needed
+             -- and extract them to dest/lib (same directory as libtorch_cuda.so)
+             downloadNvidiaLibs dest
+
+          -- Create an idempotence marker
           writeFile marker ""
           pure $ Just dest
 
@@ -160,11 +170,259 @@
     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)."
 
+-- | Downloads necessary NVIDIA wheels from PyPI and extracts libs to dest/lib
+-- Dynamically detects which library versions are needed by inspecting libtorch
+-- Uses iterative detection to find transitive dependencies
+downloadNvidiaLibs :: FilePath -> IO ()
+downloadNvidiaLibs dest = do
+  let libDir = dest </> "lib"
+
+  -- Pass 1: Detect and download direct dependencies from libtorch
+  putStrLn "Pass 1: Detecting NVIDIA library dependencies from libtorch..."
+  neededLibs <- detectNeededNvidiaLibs libDir
+
+  if null neededLibs
+    then putStrLn "No missing NVIDIA libraries detected (all bundled with libtorch)."
+    else do
+      putStrLn $ "Found " ++ show (length neededLibs) ++ " NVIDIA libraries to download:"
+      forM_ neededLibs $ \(lib, ver) ->
+        putStrLn $ "  - " ++ lib ++ " (version " ++ ver ++ ")"
+      downloadAndExtractPackages dest neededLibs
+
+      -- Pass 2: Check downloaded libraries for transitive dependencies
+      putStrLn "\nPass 2: Checking downloaded libraries for transitive dependencies..."
+      transitiveLibs <- detectNeededNvidiaLibs libDir
+
+      if null transitiveLibs
+        then putStrLn "No additional transitive dependencies found."
+        else do
+          putStrLn $ "Found " ++ show (length transitiveLibs) ++ " additional libraries to download:"
+          forM_ transitiveLibs $ \(lib, ver) ->
+            putStrLn $ "  - " ++ lib ++ " (version " ++ ver ++ ")"
+          downloadAndExtractPackages dest transitiveLibs
+
+-- | Detect which NVIDIA libraries are needed by inspecting libtorch's dependencies
+-- Returns list of (library name, major version) tuples, e.g., [("cusparse", "12"), ("cufft", "11")]
+detectNeededNvidiaLibs :: FilePath -> IO [(String, String)]
+detectNeededNvidiaLibs libDir = do
+  putStrLn $ "  Inspecting libraries in: " ++ libDir
+
+  -- Libraries we care about (including transitive dependencies)
+  let targetLibs = ["cusparse", "cufft", "curand", "cublas", "cusolver", "nvjitlink"]
+
+  -- Check libtorch libraries first
+  let libtorchFiles = ["libtorch_cuda.so", "libtorch.so"]
+
+  -- Also check any NVIDIA libraries that were already downloaded (for transitive deps)
+  libDirContents <- listDirectory libDir
+  let nvidiaLibs = filter (\f -> any (\target -> ("lib" ++ target) `isPrefixOf` f && ".so" `isInfixOf` f) targetLibs) libDirContents
+
+  let soFiles = libtorchFiles ++ nvidiaLibs
+
+  needed <- forM soFiles $ \soFile -> do
+    let fullPath = libDir </> soFile
+    exists <- doesFileExist fullPath
+    if not exists
+      then do
+        putStrLn $ "  Skipping " ++ soFile ++ " (not found)"
+        return []
+      else do
+        putStrLn $ "  Analyzing dependencies of " ++ soFile ++ "..."
+        deps <- extractNvidiaDeps fullPath targetLibs
+        forM_ deps $ \(lib, ver) ->
+          putStrLn $ "    Found dependency: lib" ++ lib ++ ".so." ++ ver
+        return deps
+
+  -- Flatten results and remove duplicates
+  let allNeeded = concat needed
+      uniqueNeeded = nubBy (\(a,_) (b,_) -> a == b) allNeeded
+
+  putStrLn $ "  Total unique NVIDIA dependencies found: " ++ show (length uniqueNeeded)
+
+  -- Filter out libraries that are already bundled
+  filterM (notBundled libDir) uniqueNeeded
+
+-- | Check if a library is already bundled in libtorch
+notBundled :: FilePath -> (String, String) -> IO Bool
+notBundled libDir (libName, _ver) = do
+  files <- listDirectory libDir
+  -- Match exactly "libNAME.so" or "libNAME-*.so" to avoid false matches
+  -- (e.g., "libcusparse" should not match "libcusparseLt")
+  let pattern = "lib" ++ libName
+  let bundled = any (\f ->
+        (f == pattern ++ ".so" ||
+         (pattern ++ ".so.") `isPrefixOf` f ||
+         (pattern ++ "-") `isPrefixOf` f && ".so" `isInfixOf` f)) files
+  return (not bundled)
+
+-- | Extract NVIDIA library dependencies using readelf
+extractNvidiaDeps :: FilePath -> [String] -> IO [(String, String)]
+extractNvidiaDeps soFile targetLibs = do
+  -- Try readelf first
+  result <- tryReadelf soFile
+  case result of
+    Just output -> return $ parseNvidiaDeps output targetLibs
+    Nothing -> do
+      -- Fallback: try objdump
+      result2 <- tryObjdump soFile
+      case result2 of
+        Just output -> return $ parseNvidiaDeps output targetLibs
+        Nothing -> return []
+
+tryReadelf :: FilePath -> IO (Maybe String)
+tryReadelf soFile = do
+  result <- try $ readProcess "readelf" ["-d", soFile] ""
+  case result of
+    Right output -> return (Just output)
+    Left (_ :: IOException) -> return Nothing
+
+tryObjdump :: FilePath -> IO (Maybe String)
+tryObjdump soFile = do
+  result <- try $ readProcess "objdump" ["-p", soFile] ""
+  case result of
+    Right output -> return (Just output)
+    Left (_ :: IOException) -> return Nothing
+
+-- | Parse readelf/objdump output to extract NVIDIA library dependencies
+-- Example: "libcusparse.so.12" -> ("cusparse", "12")
+-- Case-insensitive matching to handle libraries like "libnvJitLink" vs "libnvjitlink"
+parseNvidiaDeps :: String -> [String] -> [(String, String)]
+parseNvidiaDeps output targetLibs =
+  [ (lib, ver)
+  | line <- lines output
+  , lib <- targetLibs
+  , let pattern = "lib" ++ lib ++ ".so."
+  , let lineLower = map toLower line
+  , let patternLower = map toLower pattern
+  , patternLower `isInfixOf` lineLower
+  , let ver = extractVersion line lib
+  , not (null ver)
+  ]
+
+extractVersion :: String -> String -> String
+extractVersion line libName =
+  case findSubstring pattern line of
+    Just idx ->
+      let afterPattern = drop (idx + length pattern) line
+      in takeWhile isDigit afterPattern
+    Nothing -> ""
+  where
+    pattern = "lib" ++ libName ++ ".so."
+    isDigit c = c >= '0' && c <= '9'
+
+    -- Case-insensitive substring search
+    findSubstring :: String -> String -> Maybe Int
+    findSubstring needle haystack = go 0 haystack
+      where
+        needleLower = map toLower needle
+        go _ [] = Nothing
+        go idx str
+          | length str >= length needle &&
+            map toLower (take (length needle) str) == needleLower = Just idx
+          | otherwise = go (idx + 1) (tail str)
+
+downloadAndExtractPackages :: FilePath -> [(String, String)] -> IO ()
+downloadAndExtractPackages dest neededLibs =
+  withSystemTempDirectory "nvidia-libs-download" $ \tmpDir -> do
+    forM_ neededLibs $ \(libName, majorVer) -> do
+      -- Special case: nvidia-cufft-cu12 provides .so.11, but we need .so.12
+      -- In this case, use the generic nvidia-cufft package which has 12.x
+      let useGenericFirst = (libName == "cufft" && majorVer == "12")
+
+      -- Map major version to PyPI suffix (12 -> cu12, 11 -> cu11, etc.)
+      let pypiSuffix = "cu" ++ majorVer
+      let pkgWithSuffix = "nvidia-" ++ libName ++ "-" ++ pypiSuffix
+      let pkgGeneric = "nvidia-" ++ libName
+
+      -- Try generic first for known mismatches, otherwise try CUDA-specific first
+      (finalPkg, finalUrl) <- if useGenericFirst
+        then do
+          putStrLn $ "Fetching metadata for " ++ pkgGeneric ++ " (known version mismatch)..."
+          mUrl <- getPyPiWheelUrl pkgGeneric
+          case mUrl of
+            Just url -> return (pkgGeneric, Just url)
+            Nothing -> do
+              putStrLn $ "  Not found, trying: " ++ pkgWithSuffix ++ "..."
+              url <- getPyPiWheelUrl pkgWithSuffix
+              return (pkgWithSuffix, url)
+        else do
+          putStrLn $ "Fetching metadata for " ++ pkgWithSuffix ++ "..."
+          mUrl <- getPyPiWheelUrl pkgWithSuffix
+          -- Fallback to generic package name if versioned one not found
+          -- (e.g., nvidia-curand instead of nvidia-curand-cu10)
+          case mUrl of
+            Just url -> return (pkgWithSuffix, Just url)
+            Nothing -> do
+              putStrLn $ "  Not found, trying fallback: " ++ pkgGeneric ++ "..."
+              url <- getPyPiWheelUrl pkgGeneric
+              return (pkgGeneric, url)
+
+      case finalUrl of
+        Nothing -> putStrLn $ "Warning: Could not find manylinux wheel for " ++ finalPkg
+        Just url -> do
+          let fileName = takeFileName url
+          let downloadPath = tmpDir </> fileName
+          putStrLn $ "Downloading " ++ fileName ++ "..."
+          
+          request  <- parseRequest url
+          response <- httpLBS request
+          LBS.writeFile downloadPath (getResponseBody response)
+          
+          putStrLn $ "Extracting libraries from " ++ fileName ++ "..."
+          archive <- toArchive <$> LBS.readFile downloadPath
+          
+          -- Iterate over entries and extract only .so files to dest/lib
+          let libDest = dest </> "lib"
+          createDirectoryIfMissing True libDest
+          
+          forM_ (zEntries archive) $ \entry -> do
+            let path = eRelativePath entry
+                isSharedObj = ".so" `isInfixOf` path
+                -- Wheels put libs in nvidia/<pkg>/lib/ or lib/
+                isLibDir = "lib/" `isInfixOf` path 
+            
+            when (isSharedObj && isLibDir) $ do
+              let entryFileName = takeFileName path
+              let targetPath = libDest </> entryFileName
+              putStrLn $ "  Extracting: " ++ entryFileName ++ " -> " ++ targetPath
+              -- Write the file
+              let entryData = fromEntry entry
+              LBS.writeFile targetPath entryData
+              -- Set readable and executable permissions for shared libraries
+              -- Using only owner permissions (older directory library compatibility)
+              let perms = foldl (flip ($)) emptyPermissions
+                          [ setOwnerReadable True
+                          , setOwnerExecutable True
+                          ]
+              setPermissions targetPath perms
+
+-- | Quick and dirty PyPI JSON parser to find manylinux x86_64 url
+-- Avoids adding Aeson dependency to Setup.hs
+getPyPiWheelUrl :: String -> IO (Maybe String)
+getPyPiWheelUrl pkg = do
+  let jsonUrl = "https://pypi.org/pypi/" ++ pkg ++ "/json"
+  request <- parseRequest jsonUrl
+  response <- httpLBS request
+  let body = getResponseBody response
+  
+  -- We look for the "url" field inside an object that also has "manylinux" and "x86_64" in the filename
+  -- This is a heuristic search on the raw JSON string
+  let urls = extractUrls body
+  return $ listToMaybe [ u | u <- urls, "manylinux" `isInfixOf` u, "x86_64" `isInfixOf` u ]
+
+-- | Helper to extract all strings that look like URLs from JSON
+extractUrls :: LBS.ByteString -> [String]
+extractUrls content = 
+  let parts = LBSC.split '"' content
+      -- Filter for https URLs
+  in [ LBSC.unpack p | p <- parts, "https://" `LBSC.isPrefixOf` p, ".whl" `LBSC.isSuffixOf` p ]
+
 copyTree :: FilePath -> FilePath -> IO ()
 copyTree src dest = do
   createDirectoryIfMissing True dest
@@ -206,5 +464,6 @@
 addRPath :: FilePath -> ProgramDb -> ProgramDb
 addRPath libDir progDb =
   userSpecifyArgs (programName ldProgram)
-  ["-Wl,-rpath," ++ libDir]
+  [ "-Wl,-rpath," ++ libDir          -- for runtime library search
+  ]
   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.10
+version:             2.0.2.0
 -- 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.
@@ -12,6 +12,7 @@
 category:            Codegen
 build-type:          Custom
 extra-source-files:  README.md
+                     CHANGELOG.md
 
 custom-setup
   setup-depends:
diff --git a/src/Torch/Internal/GC.hs b/src/Torch/Internal/GC.hs
--- a/src/Torch/Internal/GC.hs
+++ b/src/Torch/Internal/GC.hs
@@ -96,7 +96,7 @@
 #ifdef darwin_HOST_OS
     msgOutOfMemory = "MPS backend out of memory"
 #else
-    msgOutOfMemory = "Exception: CUDA out of memory."
+    msgOutOfMemory = "CUDA out of memory."
 #endif
 {-# INLINE retryWithGC' #-}
 
diff --git a/src/Torch/Internal/Managed/Optim.hs b/src/Torch/Internal/Managed/Optim.hs
--- a/src/Torch/Internal/Managed/Optim.hs
+++ b/src/Torch/Internal/Managed/Optim.hs
@@ -148,3 +148,30 @@
 
 load :: ForeignPtr Optimizer -> ForeignPtr StdString -> IO ()
 load = _cast2 Unmanaged.load
+
+adamwWithParamGroups
+  :: CDouble
+  -> CDouble
+  -> CDouble
+  -> CDouble
+  -> CDouble
+  -> CBool
+  -> ForeignPtr TensorList
+  -> ForeignPtr TensorList
+  -> IO (ForeignPtr Optimizer)
+adamwWithParamGroups = _cast8 Unmanaged.adamwWithParamGroups
+
+getAllParams :: ForeignPtr Optimizer -> IO (ForeignPtr TensorList)
+getAllParams = _cast1 Unmanaged.getAllParams
+
+stepOnly :: ForeignPtr Optimizer -> IO ()
+stepOnly = _cast1 Unmanaged.stepOnly
+
+zeroGrad :: ForeignPtr Optimizer -> IO ()
+zeroGrad = _cast1 Unmanaged.zeroGrad
+
+setParamGrads :: ForeignPtr Optimizer -> ForeignPtr TensorList -> IO ()
+setParamGrads = _cast2 Unmanaged.setParamGrads
+
+setLr :: ForeignPtr Optimizer -> CDouble -> IO ()
+setLr = _cast2 Unmanaged.setLr
diff --git a/src/Torch/Internal/Unmanaged/Optim.hs b/src/Torch/Internal/Unmanaged/Optim.hs
--- a/src/Torch/Internal/Unmanaged/Optim.hs
+++ b/src/Torch/Internal/Unmanaged/Optim.hs
@@ -261,10 +261,17 @@
     return dynamic_cast<torch::optim::Optimizer*>(optimizer);
   }|]
 
-getParams :: Ptr Optimizer -> IO (Ptr TensorList) 
+getParams :: Ptr Optimizer -> IO (Ptr TensorList)
 getParams optimizer =
   [C.throwBlock| std::vector<at::Tensor>* {
-    return new std::vector<at::Tensor>($(torch::optim::Optimizer* optimizer)->param_groups().at(0).params());
+    auto optimizer = $(torch::optim::Optimizer* optimizer);
+    std::vector<at::Tensor>* result = new std::vector<at::Tensor>();
+    for(auto& group : optimizer->param_groups()){
+      for(auto& p : group.params()){
+        result->push_back(p);
+      }
+    }
+    return result;
   }|]
 
 step :: Ptr Optimizer -> (Ptr TensorList -> IO (Ptr Tensor)) -> IO (Ptr Tensor)
@@ -280,7 +287,13 @@
         auto func = (Func)tfunc;
         auto v = optimizer->step([&]{
           optimizer->zero_grad();
-          auto loss = func(&(optimizer->param_groups().at(0).params()));
+          std::vector<at::Tensor> all_params;
+          for(auto& group : optimizer->param_groups()){
+            for(auto& p : group.params()){
+              all_params.push_back(p);
+            }
+          }
+          auto loss = func(&all_params);
           loss->backward();
           return *loss;
         });
@@ -304,7 +317,13 @@
         auto func = (Func)tfunc;
         auto v = optimizer->step([&]{
           optimizer->zero_grad();
-          auto lossWithGenerator = func(&(optimizer->param_groups().at(0).params()),&generator);
+          std::vector<at::Tensor> all_params;
+          for(auto& group : optimizer->param_groups()){
+            for(auto& p : group.params()){
+              all_params.push_back(p);
+            }
+          }
+          auto lossWithGenerator = func(&all_params,&generator);
           auto loss = std::get<0>(*lossWithGenerator);
           generator = std::get<1>(*lossWithGenerator);
           loss.backward();
@@ -326,7 +345,13 @@
     optimizer->zero_grad();
     loss->backward();
     optimizer->step();
-    return new std::vector<at::Tensor>(optimizer->param_groups().at(0).params());
+    std::vector<at::Tensor>* result = new std::vector<at::Tensor>();
+    for(auto& group : optimizer->param_groups()){
+      for(auto& p : group.params()){
+        result->push_back(p);
+      }
+    }
+    return result;
   }|]
 
 save :: Ptr Optimizer -> Ptr StdString -> IO ()
@@ -341,4 +366,95 @@
   [C.throwBlock| void {
     std::ifstream input(*$(std::string* filename));
     torch::load(*$(torch::optim::Optimizer* optimizer),input);
+  }|]
+
+adamwWithParamGroups
+  :: CDouble
+  -> CDouble
+  -> CDouble
+  -> CDouble
+  -> CDouble
+  -> CBool
+  -> Ptr TensorList
+  -> Ptr TensorList
+  -> IO (Ptr Optimizer)
+adamwWithParamGroups adamLr adamBetas0 adamBetas1 adamEps adamWeightDecay adamAmsgrad decayParams noDecayParams =
+  [C.throwBlock| torch::optim::Optimizer* {
+    std::vector<at::Tensor>* decay_params = $(std::vector<at::Tensor>* decayParams);
+    std::vector<at::Tensor>* no_decay_params = $(std::vector<at::Tensor>* noDecayParams);
+    std::vector<at::Tensor> dp;
+    for(int i=0;i<decay_params->size();i++){
+      dp.push_back((*decay_params)[i].detach().set_requires_grad(true));
+    }
+    std::vector<at::Tensor> ndp;
+    for(int i=0;i<no_decay_params->size();i++){
+      ndp.push_back((*no_decay_params)[i].detach().set_requires_grad(true));
+    }
+    auto options_decay = torch::optim::AdamWOptions()
+      .lr($(double adamLr))
+      .betas(std::make_tuple($(double adamBetas0),$(double adamBetas1)))
+      .eps($(double adamEps))
+      .weight_decay($(double adamWeightDecay))
+      .amsgrad($(bool adamAmsgrad));
+    auto options_no_decay = torch::optim::AdamWOptions()
+      .lr($(double adamLr))
+      .betas(std::make_tuple($(double adamBetas0),$(double adamBetas1)))
+      .eps($(double adamEps))
+      .weight_decay(0.0)
+      .amsgrad($(bool adamAmsgrad));
+    std::vector<torch::optim::OptimizerParamGroup> param_groups;
+    param_groups.emplace_back(torch::optim::OptimizerParamGroup(ndp, std::make_unique<torch::optim::AdamWOptions>(options_no_decay)));
+    param_groups.emplace_back(torch::optim::OptimizerParamGroup(dp, std::make_unique<torch::optim::AdamWOptions>(options_decay)));
+    torch::optim::AdamW* optimizer = new torch::optim::AdamW(param_groups);
+    optimizer->zero_grad();
+    return dynamic_cast<torch::optim::Optimizer*>(optimizer);
+  }|]
+
+getAllParams :: Ptr Optimizer -> IO (Ptr TensorList)
+getAllParams optimizer =
+  [C.throwBlock| std::vector<at::Tensor>* {
+    auto optimizer = $(torch::optim::Optimizer* optimizer);
+    std::vector<at::Tensor>* result = new std::vector<at::Tensor>();
+    for(auto& group : optimizer->param_groups()){
+      for(auto& p : group.params()){
+        result->push_back(p);
+      }
+    }
+    return result;
+  }|]
+
+stepOnly :: Ptr Optimizer -> IO ()
+stepOnly optimizer =
+  [C.throwBlock| void {
+    $(torch::optim::Optimizer* optimizer)->step();
+  }|]
+
+zeroGrad :: Ptr Optimizer -> IO ()
+zeroGrad optimizer =
+  [C.throwBlock| void {
+    $(torch::optim::Optimizer* optimizer)->zero_grad();
+  }|]
+
+setParamGrads :: Ptr Optimizer -> Ptr TensorList -> IO ()
+setParamGrads optimizer grads =
+  [C.throwBlock| void {
+    auto optimizer = $(torch::optim::Optimizer* optimizer);
+    auto grads = $(std::vector<at::Tensor>* grads);
+    int idx = 0;
+    for(auto& group : optimizer->param_groups()){
+      for(auto& p : group.params()){
+        p.mutable_grad() = (*grads)[idx];
+        idx++;
+      }
+    }
+  }|]
+
+setLr :: Ptr Optimizer -> CDouble -> IO ()
+setLr optimizer newLr =
+  [C.throwBlock| void {
+    auto optimizer = $(torch::optim::Optimizer* optimizer);
+    for(auto& group : optimizer->param_groups()){
+      auto& options = static_cast<torch::optim::AdamWOptions&>(group.options());
+      options.lr($(double newLr));
+    }
   }|]