diff --git a/.github/workflows/engine.yml b/.github/workflows/engine.yml
index 60a5e3f9..ac2392d3 100644
--- a/.github/workflows/engine.yml
+++ b/.github/workflows/engine.yml
@@ -29,7 +29,7 @@ jobs:
           persist-credentials: false
       - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
         with:
-          toolchain: nightly-2024-07-18
+          toolchain: nightly-2025-02-17
           override: true
           components: rustfmt, clippy
       - uses: jetli/wasm-pack-action@0d096b08b4e5a7de8c28de67e11e945404e9eefa  #v0.4.0
@@ -40,7 +40,7 @@ jobs:
         run: rustup target add wasm32-unknown-unknown
 
       - name: Add rust-src
-        run: rustup component add rust-src --toolchain nightly-2024-07-18-x86_64-unknown-linux-gnu
+        run: rustup component add rust-src --toolchain nightly-2025-02-17-x86_64-unknown-linux-gnu
       - name: Install binaryen
         run: |
           set -e
diff --git a/.github/workflows/large-tests.yml b/.github/workflows/large-tests.yml
index 2aac6a0f..e57de2ba 100644
--- a/.github/workflows/large-tests.yml
+++ b/.github/workflows/large-tests.yml
@@ -15,7 +15,7 @@ jobs:
           persist-credentials: false
       - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
         with:
-          toolchain: nightly-2024-07-18
+          toolchain: nightly-2025-02-17
           override: true
           components: rustfmt, clippy
       - name: nanoGPT Mock
diff --git a/.github/workflows/pypi.yml b/.github/workflows/pypi.yml
index bba90268..ab830e08 100644
--- a/.github/workflows/pypi.yml
+++ b/.github/workflows/pypi.yml
@@ -50,7 +50,7 @@ jobs:
 
       - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
         with:
-          toolchain: nightly-2024-07-18
+          toolchain: nightly-2025-02-17
           override: true
           components: rustfmt, clippy
 
@@ -115,7 +115,7 @@ jobs:
 
       - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
         with:
-          toolchain: nightly-2024-07-18
+          toolchain: nightly-2025-02-17
           override: true
           components: rustfmt, clippy
 
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
index 5002b313..7b25978d 100644
--- a/.github/workflows/release.yml
+++ b/.github/workflows/release.yml
@@ -51,7 +51,7 @@ jobs:
     steps:
       - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
         with:
-          toolchain: nightly-2024-07-18
+          toolchain: nightly-2025-02-17
           override: true
           components: rustfmt, clippy
       - name: Checkout repo
@@ -119,27 +119,27 @@ jobs:
         include:
           - build: windows-msvc
             os: windows-latest
-            rust: nightly-2024-07-18
+            rust: nightly-2025-02-17
             target: x86_64-pc-windows-msvc
           - build: macos
             os: macos-13
-            rust: nightly-2024-07-18
+            rust: nightly-2025-02-17
             target: x86_64-apple-darwin
           - build: macos-aarch64
             os: macos-13
-            rust: nightly-2024-07-18
+            rust: nightly-2025-02-17
             target: aarch64-apple-darwin
           - build: linux-musl
             os: ubuntu-22.04
-            rust: nightly-2024-07-18
+            rust: nightly-2025-02-17
             target: x86_64-unknown-linux-musl
           - build: linux-gnu
             os: ubuntu-22.04
-            rust: nightly-2024-07-18
+            rust: nightly-2025-02-17
             target: x86_64-unknown-linux-gnu
           - build: linux-aarch64
             os: ubuntu-22.04
-            rust: nightly-2024-07-18
+            rust: nightly-2025-02-17
             target: aarch64-unknown-linux-gnu
 
     steps:
diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml
index afd1eaa8..4782cd7c 100644
--- a/.github/workflows/rust.yml
+++ b/.github/workflows/rust.yml
@@ -30,7 +30,7 @@ jobs:
           persist-credentials: false
       - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
         with:
-          toolchain: nightly-2024-07-18
+          toolchain: nightly-2025-02-17
           override: true
           components: rustfmt, clippy
       - uses: baptiste0928/cargo-install@v1
@@ -50,7 +50,7 @@ jobs:
           persist-credentials: false
       - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
         with:
-          toolchain: nightly-2024-07-18
+          toolchain: nightly-2025-02-17
           override: true
           components: rustfmt, clippy
       - name: Build
@@ -66,7 +66,7 @@ jobs:
           persist-credentials: false
       - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
         with:
-          toolchain: nightly-2024-07-18
+          toolchain: nightly-2025-02-17
           override: true
           components: rustfmt, clippy
       - name: Docs
@@ -82,7 +82,7 @@ jobs:
           persist-credentials: false
       - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
         with:
-          toolchain: nightly-2024-07-18
+          toolchain: nightly-2025-02-17
           override: true
           components: rustfmt, clippy
       - uses: baptiste0928/cargo-install@91c5da15570085bcde6f4d7aed98cb82d6769fd3  #v3.3.0
@@ -107,7 +107,7 @@ jobs:
   #   persist-credentials: false
   #     - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
   #       with:
-  #         toolchain: nightly-2024-07-18
+  #         toolchain: nightly-2025-02-17
   #         override: true
   #         components: rustfmt, clippy
   #     - uses: baptiste0928/cargo-install@91c5da15570085bcde6f4d7aed98cb82d6769fd3  #v3.3.0
@@ -117,10 +117,6 @@ jobs:
   #     - uses: mwilliamson/setup-wasmtime-action@bf814d7d8fc3c3a77dfe114bd9fb8a2c575f6ad6  #v2.0.0
   #       with:
   #         wasmtime-version: "3.0.1"
-  #     - name: Install wasm32-wasi
-  #       run: rustup target add wasm32-wasi
-  #     - name: Install cargo-wasi
-  #       run: cargo install cargo-wasi
   #     # - name: Matmul overflow (wasi)
   #     #   run: cargo wasi test matmul_col_ultra_overflow -- --include-ignored --nocapture
   #     # - name: Conv overflow (wasi)
@@ -144,7 +140,7 @@ jobs:
           persist-credentials: false
       - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
         with:
-          toolchain: nightly-2024-07-18
+          toolchain: nightly-2025-02-17
           override: true
           components: rustfmt, clippy
       - uses: baptiste0928/cargo-install@91c5da15570085bcde6f4d7aed98cb82d6769fd3  #v3.3.0
@@ -154,10 +150,6 @@ jobs:
       - uses: mwilliamson/setup-wasmtime-action@bf814d7d8fc3c3a77dfe114bd9fb8a2c575f6ad6  #v2.0.0
         with:
           wasmtime-version: "3.0.1"
-      - name: Install wasm32-wasi
-        run: rustup target add wasm32-wasi
-      - name: Install cargo-wasi
-        run: cargo install cargo-wasi
       # - name: Matmul overflow (wasi)
       #   run: cargo wasi test matmul_col_ultra_overflow -- --include-ignored --nocapture
       # - name: Conv overflow (wasi)
@@ -181,7 +173,7 @@ jobs:
           persist-credentials: false
       - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
         with:
-          toolchain: nightly-2024-07-18
+          toolchain: nightly-2025-02-17
           override: true
           components: rustfmt, clippy
       - uses: baptiste0928/cargo-install@91c5da15570085bcde6f4d7aed98cb82d6769fd3  #v3.3.0
@@ -191,10 +183,6 @@ jobs:
       - uses: mwilliamson/setup-wasmtime-action@bf814d7d8fc3c3a77dfe114bd9fb8a2c575f6ad6  #v2.0.0
         with:
           wasmtime-version: "3.0.1"
-      - name: Install wasm32-wasi
-        run: rustup target add wasm32-wasi
-      - name: Install cargo-wasi
-        run: cargo install cargo-wasi
       # - name: Matmul overflow (wasi)
       #   run: cargo wasi test matmul_col_ultra_overflow -- --include-ignored --nocapture
       # - name: Conv overflow (wasi)
@@ -218,7 +206,7 @@ jobs:
           persist-credentials: false
       - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
         with:
-          toolchain: nightly-2024-07-18
+          toolchain: nightly-2025-02-17
           override: true
           components: rustfmt, clippy
       - uses: baptiste0928/cargo-install@91c5da15570085bcde6f4d7aed98cb82d6769fd3  #v3.3.0
@@ -238,7 +226,7 @@ jobs:
           persist-credentials: false
       - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
         with:
-          toolchain: nightly-2024-07-18
+          toolchain: nightly-2025-02-17
           override: true
           components: rustfmt, clippy
       - uses: jetli/wasm-pack-action@0d096b08b4e5a7de8c28de67e11e945404e9eefa  #v0.4.0
@@ -251,7 +239,7 @@ jobs:
       - name: Install wasm32-unknown-unknown
         run: rustup target add wasm32-unknown-unknown
       - name: Add rust-src
-        run: rustup component add rust-src --toolchain nightly-2024-07-18-x86_64-unknown-linux-gnu
+        run: rustup component add rust-src --toolchain nightly-2025-02-17-x86_64-unknown-linux-gnu
       - name: Run wasm verifier tests
         # on mac:
         #   AR=/opt/homebrew/opt/llvm/bin/llvm-ar CC=/opt/homebrew/opt/llvm/bin/clang wasm-pack test --firefox --headless -- -Z build-std="panic_abort,std" --features web
@@ -267,7 +255,7 @@ jobs:
           persist-credentials: false
       - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
         with:
-          toolchain: nightly-2024-07-18
+          toolchain: nightly-2025-02-17
           override: true
           components: rustfmt, clippy
       - uses: baptiste0928/cargo-install@91c5da15570085bcde6f4d7aed98cb82d6769fd3  #v3.3.0
@@ -334,7 +322,7 @@ jobs:
           persist-credentials: false
       - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
         with:
-          toolchain: nightly-2024-07-18
+          toolchain: nightly-2025-02-17
           override: true
           components: rustfmt, clippy
       - uses: baptiste0928/cargo-install@91c5da15570085bcde6f4d7aed98cb82d6769fd3  #v3.3.0
@@ -354,7 +342,7 @@ jobs:
           node-version: "18.12.1"
           cache: "pnpm"
       - name: "Add rust-src"
-        run: rustup component add rust-src --toolchain nightly-2024-07-18-x86_64-unknown-linux-gnu
+        run: rustup component add rust-src --toolchain nightly-2025-02-17-x86_64-unknown-linux-gnu
       - name: Install dependencies for js tests and in-browser-evm-verifier package
         run: |
           pnpm install --frozen-lockfile
@@ -419,7 +407,7 @@ jobs:
   #         persist-credentials: false
   #     - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
   #       with:
-  #         toolchain: nightly-2024-07-18
+  #         toolchain: nightly-2025-02-17
   #         override: true
   #         components: rustfmt, clippy
   #     - uses: jetli/wasm-pack-action@0d096b08b4e5a7de8c28de67e11e945404e9eefa  #v0.4.0
@@ -427,7 +415,7 @@ jobs:
   #         # Pin to version 0.12.1
   #         version: 'v0.12.1'
   #     - name: Add rust-src
-  #       run: rustup component add rust-src --toolchain nightly-2024-07-18
+  #       run: rustup component add rust-src --toolchain nightly-2025-02-17
   #     - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  #v4.2.2
   #       with:
   #         persist-credentials: false
@@ -453,7 +441,7 @@ jobs:
           persist-credentials: false
       - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
         with:
-          toolchain: nightly-2024-07-18
+          toolchain: nightly-2025-02-17
           override: true
           components: rustfmt, clippy
       - uses: jetli/wasm-pack-action@0d096b08b4e5a7de8c28de67e11e945404e9eefa  #v0.4.0
@@ -464,7 +452,7 @@ jobs:
         run: rustup target add wasm32-unknown-unknown
 
       - name: Add rust-src
-        run: rustup component add rust-src --toolchain nightly-2024-07-18-x86_64-unknown-linux-gnu
+        run: rustup component add rust-src --toolchain nightly-2025-02-17-x86_64-unknown-linux-gnu
       - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  #v4.2.2
         with:
           persist-credentials: false
@@ -534,11 +522,11 @@ jobs:
   #   persist-credentials: false
   #     - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
   #       with:
-  #         toolchain: nightly-2024-07-18
+  #         toolchain: nightly-2025-02-17
   #         override: true
   #         components: rustfmt, clippy
   #     - name: Add rust-src
-  #       run: rustup component add rust-src --toolchain nightly-2024-07-18-x86_64-unknown-linux-gnu
+  #       run: rustup component add rust-src --toolchain nightly-2025-02-17-x86_64-unknown-linux-gnu
   #     - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  #v4.2.2
   #     - uses: baptiste0928/cargo-install@91c5da15570085bcde6f4d7aed98cb82d6769fd3  #v3.3.0
   #       with:
@@ -572,7 +560,7 @@ jobs:
           persist-credentials: false
       - uses: dtolnay/rust-toolchain@4f94fbe7e03939b0e674bcc9ca609a16088f63ff  #nightly branch, TODO: update when required
         with:
-          toolchain: nightly-2024-07-18
+          toolchain: nightly-2025-02-17
           override: true
           components: rustfmt, clippy
       - uses: baptiste0928/cargo-install@91c5da15570085bcde6f4d7aed98cb82d6769fd3  #v3.3.0
@@ -592,7 +580,7 @@ jobs:
   #   persist-credentials: false
   #     - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
   #       with:
-  #         toolchain: nightly-2024-07-18
+  #         toolchain: nightly-2025-02-17
   #         override: true
   #         components: rustfmt, clippy
   #     - uses: baptiste0928/cargo-install@91c5da15570085bcde6f4d7aed98cb82d6769fd3  #v3.3.0
@@ -613,7 +601,7 @@ jobs:
           persist-credentials: false
       - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
         with:
-          toolchain: nightly-2024-07-18
+          toolchain: nightly-2025-02-17
           override: true
           components: rustfmt, clippy
       - uses: baptiste0928/cargo-install@91c5da15570085bcde6f4d7aed98cb82d6769fd3  #v3.3.0
@@ -634,7 +622,7 @@ jobs:
           persist-credentials: false
       - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
         with:
-          toolchain: nightly-2024-07-18
+          toolchain: nightly-2025-02-17
           override: true
           components: rustfmt, clippy
       - uses: baptiste0928/cargo-install@91c5da15570085bcde6f4d7aed98cb82d6769fd3  #v3.3.0
@@ -659,7 +647,7 @@ jobs:
           persist-credentials: false
       - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
         with:
-          toolchain: nightly-2024-07-18
+          toolchain: nightly-2025-02-17
           override: true
           components: rustfmt, clippy
       - uses: baptiste0928/cargo-install@91c5da15570085bcde6f4d7aed98cb82d6769fd3  #v3.3.0
@@ -683,7 +671,7 @@ jobs:
           python-version: "3.12"
       - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
         with:
-          toolchain: nightly-2024-07-18
+          toolchain: nightly-2025-02-17
           override: true
           components: rustfmt, clippy
       - name: Install cmake
@@ -713,7 +701,7 @@ jobs:
           python-version: "3.12"
       - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
         with:
-          toolchain: nightly-2024-07-18
+          toolchain: nightly-2025-02-17
           override: true
           components: rustfmt, clippy
       - uses: baptiste0928/cargo-install@91c5da15570085bcde6f4d7aed98cb82d6769fd3  #v3.3.0
@@ -764,7 +752,7 @@ jobs:
           python-version: "3.11"
       - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
         with:
-          toolchain: nightly-2024-07-18
+          toolchain: nightly-2025-02-17
           override: true
           components: rustfmt, clippy
       - uses: baptiste0928/cargo-install@91c5da15570085bcde6f4d7aed98cb82d6769fd3  #v3.3.0
@@ -819,7 +807,7 @@ jobs:
           persist-credentials: false
       - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
         with:
-          toolchain: nightly-2024-07-18
+          toolchain: nightly-2025-02-17
           override: true
           components: rustfmt, clippy
       - uses: baptiste0928/cargo-install@91c5da15570085bcde6f4d7aed98cb82d6769fd3  #v3.3.0
@@ -827,7 +815,7 @@ jobs:
           crate: cargo-nextest
           locked: true
       - name: Run ios tests
-        run: CARGO_BUILD_TARGET=aarch64-apple-darwin RUSTUP_TOOLCHAIN=nightly-2024-07-18-aarch64-apple-darwin cargo test --test ios_integration_tests --features ios-bindings-test --no-default-features
+        run: CARGO_BUILD_TARGET=aarch64-apple-darwin RUSTUP_TOOLCHAIN=nightly-2025-02-17-aarch64-apple-darwin cargo test --test ios_integration_tests --features ios-bindings-test --no-default-features
 
   swift-package-tests:
     permissions:
@@ -841,7 +829,7 @@ jobs:
           persist-credentials: false
       - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
         with:
-          toolchain: nightly-2024-07-18
+          toolchain: nightly-2025-02-17
           override: true
           components: rustfmt, clippy
       - name: Build EzklCoreBindings
diff --git a/.github/workflows/static-analysis.yml b/.github/workflows/static-analysis.yml
index 45fef5f3..3c3f8622 100644
--- a/.github/workflows/static-analysis.yml
+++ b/.github/workflows/static-analysis.yml
@@ -17,7 +17,7 @@ jobs:
           persist-credentials: false
       - uses: actions-rs/toolchain@b2417cde72dcf67f306c0ae8e0828a81bf0b189f  #v1.0.6
         with:
-          toolchain: nightly-2024-07-18
+          toolchain: nightly-2025-02-17
           override: true
           components: rustfmt, clippy
 
diff --git a/Cargo.toml b/Cargo.toml
index 2848956f..457d990f 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -3,7 +3,7 @@ cargo-features = ["profile-rustflags"]
 [package]
 name = "ezkl"
 version = "0.0.0"
-edition = "2021"
+edition = "2024"
 default-run = "ezkl"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
diff --git a/benches/zero_finder.rs b/benches/zero_finder.rs
index 80da6ef4..73123bfe 100644
--- a/benches/zero_finder.rs
+++ b/benches/zero_finder.rs
@@ -22,7 +22,7 @@ fn generate_test_data(size: usize, zero_probability: f64) -> Vec<ValType> {
     let mut rng = rand::thread_rng();
     (0..size)
         .map(|_i| {
-            if rng.gen::<f64>() < zero_probability {
+            if rng.r#gen::<f64>() < zero_probability {
                 ValType::Constant(F::ZERO)
             } else {
                 ValType::Constant(F::ONE) // Or some other non-zero value
diff --git a/examples/notebooks/cat_and_dog.ipynb b/examples/notebooks/cat_and_dog.ipynb
index 2598fb90..af0c5888 100644
--- a/examples/notebooks/cat_and_dog.ipynb
+++ b/examples/notebooks/cat_and_dog.ipynb
@@ -373,15 +373,19 @@
             "outputs": [],
             "source": [
                 "# Set image size.\n",
-                "IMAGE_WIDTH = 112\n",
-                "IMAGE_HEIGHT = 112\n",
+                "IMAGE_WIDTH = 64\n",
+                "IMAGE_HEIGHT = 64\n",
                 "IMAGE_SIZE=(IMAGE_WIDTH, IMAGE_HEIGHT)\n",
                 "\n",
                 "# Create training transform with TrivialAugment\n",
                 "train_transform = transforms.Compose([\n",
                 "    transforms.Resize(IMAGE_SIZE),\n",
                 "    transforms.TrivialAugmentWide(),\n",
-                "    transforms.ToTensor()])\n",
+                "     transforms.RandomHorizontalFlip(),\n",
+                "    transforms.RandomRotation(10),\n",
+                "    transforms.ColorJitter(brightness=0.2, contrast=0.2),\n",
+                "    transforms.ToTensor(),\n",
+                "    ])\n",
                 "\n",
                 "# Create testing transform (no data augmentation)\n",
                 "test_transform = transforms.Compose([\n",
@@ -424,7 +428,7 @@
         },
         {
             "cell_type": "code",
-            "execution_count": 26,
+            "execution_count": 74,
             "metadata": {},
             "outputs": [],
             "source": [
@@ -434,25 +438,55 @@
                 "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
                 "device\n",
                 "\n",
-                "# # Creating a CNN-based image classifier.\n",
+                "#  Improved CNN-based image classifier\n",
                 "class ImageClassifier(nn.Module):\n",
                 "    def __init__(self):\n",
                 "        super().__init__()\n",
+                "        \n",
+                "        # First convolutional block with batch normalization and LeakyReLU\n",
                 "        self.conv_layer_1 = nn.Sequential(\n",
-                "          nn.Conv2d(3, 4, 3, padding=1),\n",
-                "          nn.ReLU(),\n",
-                "          nn.MaxPool2d(2))\n",
+                "            nn.Conv2d(3, 6, 3, padding=1),  # Moderate increase from 4 to 6\n",
+                "            nn.BatchNorm2d(6),\n",
+                "            nn.LeakyReLU(0.1),\n",
+                "            nn.MaxPool2d(2)\n",
+                "        )\n",
+                "        \n",
+                "        # Second convolutional block with batch normalization and LeakyReLU\n",
                 "        self.conv_layer_2 = nn.Sequential(\n",
-                "          nn.Conv2d(4, 4, 3, padding=1),\n",
-                "          nn.ReLU(),\n",
-                "          nn.MaxPool2d(2))\n",
+                "            nn.Conv2d(6, 8, 3, padding=1),  # Moderate increase from 4 to 8\n",
+                "            nn.BatchNorm2d(8),\n",
+                "            nn.LeakyReLU(0.1),\n",
+                "            nn.MaxPool2d(2)\n",
+                "        )\n",
+                "        \n",
+                "        # For a 64x64 input, after 2 MaxPool2d(2) layers, the spatial dimensions are 16x16\n",
+                "        # With 8 channels, the flattened size is 16*16*8 = 2048\n",
+                "        \n",
+                "        # Classifier with dropout\n",
                 "        self.classifier = nn.Sequential(\n",
-                "          nn.Flatten(),\n",
-                "          nn.Linear(in_features=3136, out_features=2))\n",
-                "    def forward(self, x: torch.Tensor):\n",
+                "            nn.Flatten(),\n",
+                "            nn.Dropout(0.25),  # Add dropout for regularization\n",
+                "            nn.Linear(in_features=16*16*8, out_features=2)\n",
+                "        )\n",
+                "        \n",
+                "        # For residual connection\n",
+                "        self.downsample = nn.Sequential(\n",
+                "            nn.Conv2d(3, 8, 1, stride=4),  # Match spatial dimensions (64x64 -> 16x16)\n",
+                "            nn.BatchNorm2d(8)\n",
+                "        )\n",
+                "        \n",
+                "    def forward(self, x):\n",
+                "        # Save input for residual connection\n",
+                "        identity = self.downsample(x)\n",
+                "        \n",
                 "        x = self.conv_layer_1(x)\n",
                 "        x = self.conv_layer_2(x)\n",
+                "        \n",
+                "        # Add residual connection\n",
+                "        x = x + identity\n",
+                "        \n",
                 "        x = self.classifier(x)\n",
+                "        \n",
                 "        return x\n",
                 "# Instantiate an object.\n",
                 "model = ImageClassifier().to(device)\n"
@@ -526,7 +560,7 @@
         },
         {
             "cell_type": "code",
-            "execution_count": 29,
+            "execution_count": 77,
             "metadata": {},
             "outputs": [],
             "source": [
@@ -662,15 +696,17 @@
                 "torch.manual_seed(42) \n",
                 "torch.cuda.manual_seed(42)\n",
                 "\n",
-                "# Set number of epochs\n",
+                "# Set number of epochs (change to 1000 for better results)\n",
                 "NUM_EPOCHS = 25\n",
+                "# NUM_EPOCHS = 1000\n",
+                "\n",
                 "\n",
                 "# Setup loss function and optimizer\n",
                 "loss_fn = nn.CrossEntropyLoss()\n",
                 "optimizer = torch.optim.Adam(params=model.parameters(), lr=1e-3)\n",
                 "\n",
                 "# Start the timer\n",
-                "from timeit import default_timer as timer \n",
+                "from timeit import default_timer as timer\n",
                 "start_time = timer()\n",
                 "\n",
                 "# Train model_0\n",
@@ -695,7 +731,7 @@
         },
         {
             "cell_type": "code",
-            "execution_count": 78,
+            "execution_count": 94,
             "metadata": {},
             "outputs": [],
             "source": [
@@ -830,7 +866,7 @@
         },
         {
             "cell_type": "code",
-            "execution_count": 86,
+            "execution_count": 98,
             "metadata": {},
             "outputs": [],
             "source": [
@@ -1016,7 +1052,6 @@
             "source": [
                 "import os\n",
                 "\n",
-                "\n",
                 "res = await ezkl.create_evm_verifier()\n",
                 "\n",
                 "assert res == True"
@@ -1102,7 +1137,7 @@
             "name": "python",
             "nbconvert_exporter": "python",
             "pygments_lexer": "ipython3",
-            "version": "3.12.7"
+            "version": "3.12.9"
         }
     },
     "nbformat": 4,
diff --git a/rust-toolchain b/rust-toolchain
index 9c738d7f..45b6b172 100644
--- a/rust-toolchain
+++ b/rust-toolchain
@@ -1,3 +1,3 @@
 [toolchain]
-channel = "nightly-2024-07-18"
+channel = "nightly-2025-02-17"
 components = ["rustfmt", "clippy"]
diff --git a/src/commands.rs b/src/commands.rs
index 2d2e22ae..fe9d1ca3 100644
--- a/src/commands.rs
+++ b/src/commands.rs
@@ -1,6 +1,6 @@
 use alloy::primitives::Address as H160;
 use clap::{Command, Parser, Subcommand};
-use clap_complete::{generate, Generator, Shell};
+use clap_complete::{Generator, Shell, generate};
 #[cfg(feature = "python-bindings")]
 use pyo3::{conversion::FromPyObject, exceptions::PyValueError, prelude::*};
 use serde::{Deserialize, Serialize};
@@ -8,7 +8,7 @@ use std::path::PathBuf;
 use std::str::FromStr;
 use tosubcommand::{ToFlags, ToSubcommand};
 
-use crate::{pfsys::ProofType, Commitments, RunArgs};
+use crate::{Commitments, RunArgs, pfsys::ProofType};
 
 use crate::circuit::CheckMode;
 use crate::graph::TestDataSource;
@@ -360,8 +360,13 @@ pub fn get_styles() -> clap::builder::Styles {
 }
 
 /// Print completions for the given generator
-pub fn print_completions<G: Generator>(gen: G, cmd: &mut Command) {
-    generate(gen, cmd, cmd.get_name().to_string(), &mut std::io::stdout());
+pub fn print_completions<G: Generator>(r#gen: G, cmd: &mut Command) {
+    generate(
+        r#gen,
+        cmd,
+        cmd.get_name().to_string(),
+        &mut std::io::stdout(),
+    );
 }
 
 #[allow(missing_docs)]
diff --git a/src/execute.rs b/src/execute.rs
index 6c0a78f8..2b7964fe 100644
--- a/src/execute.rs
+++ b/src/execute.rs
@@ -1,5 +1,6 @@
-use crate::circuit::region::RegionSettings;
+use crate::EZKL_BUF_CAPACITY;
 use crate::circuit::CheckMode;
+use crate::circuit::region::RegionSettings;
 use crate::commands::CalibrationTarget;
 use crate::eth::{
     deploy_contract_via_solidity, deploy_da_verifier_via_solidity, fix_da_multi_sol,
@@ -12,21 +13,21 @@ use crate::graph::{GraphCircuit, GraphSettings, GraphWitness, Model};
 use crate::graph::{TestDataSource, TestSources};
 use crate::pfsys::evm::aggregation_kzg::{AggregationCircuit, PoseidonTranscript};
 use crate::pfsys::{
-    create_keys, load_pk, load_vk, save_params, save_pk, Snark, StrategyType, TranscriptType,
+    ProofSplitCommit, create_proof_circuit, swap_proof_commitments_polycommit, verify_proof_circuit,
 };
 use crate::pfsys::{
-    create_proof_circuit, swap_proof_commitments_polycommit, verify_proof_circuit, ProofSplitCommit,
+    Snark, StrategyType, TranscriptType, create_keys, load_pk, load_vk, save_params, save_pk,
 };
 use crate::pfsys::{save_vk, srs::*};
 use crate::tensor::TensorError;
-use crate::EZKL_BUF_CAPACITY;
-use crate::{commands::*, EZKLError};
 use crate::{Commitments, RunArgs};
+use crate::{EZKLError, commands::*};
 use colored::Colorize;
 #[cfg(unix)]
 use gag::Gag;
 use halo2_proofs::dev::VerifyFailure;
 use halo2_proofs::plonk::{self, Circuit};
+use halo2_proofs::poly::VerificationStrategy;
 use halo2_proofs::poly::commitment::{CommitmentScheme, Params};
 use halo2_proofs::poly::commitment::{ParamsProver, Verifier};
 use halo2_proofs::poly::ipa::commitment::{IPACommitmentScheme, ParamsIPA};
@@ -39,7 +40,6 @@ use halo2_proofs::poly::kzg::strategy::AccumulatorStrategy as KZGAccumulatorStra
 use halo2_proofs::poly::kzg::{
     commitment::ParamsKZG, strategy::SingleStrategy as KZGSingleStrategy,
 };
-use halo2_proofs::poly::VerificationStrategy;
 use halo2_proofs::transcript::{EncodedChallenge, TranscriptReadBuffer};
 use halo2_solidity_verifier;
 use halo2curves::bn256::{Bn256, Fr, G1Affine};
@@ -50,12 +50,12 @@ use instant::Instant;
 use itertools::Itertools;
 use log::debug;
 use log::{info, trace, warn};
-use serde::de::DeserializeOwned;
 use serde::Serialize;
+use serde::de::DeserializeOwned;
 use snark_verifier::loader::native::NativeLoader;
+use snark_verifier::system::halo2::Config;
 use snark_verifier::system::halo2::compile;
 use snark_verifier::system::halo2::transcript::evm::EvmTranscript;
-use snark_verifier::system::halo2::Config;
 use std::fs::File;
 use std::io::BufWriter;
 use std::io::{Cursor, Write};
@@ -516,7 +516,9 @@ fn update_ezkl_binary(version: &Option<String>) -> Result<String, EZKLError> {
         .status()
         .is_err()
     {
-        log::warn!("bash is not installed on this system, trying to run the install script with sh (may fail)");
+        log::warn!(
+            "bash is not installed on this system, trying to run the install script with sh (may fail)"
+        );
         "sh"
     } else {
         "bash"
@@ -876,7 +878,7 @@ pub(crate) fn gen_random_data(
 
         let mut tensor = TractTensor::zero::<f32>(sizes).unwrap();
         let slice = tensor.as_slice_mut::<f32>().unwrap();
-        slice.iter_mut().for_each(|x| *x = rng.gen());
+        slice.iter_mut().for_each(|x| *x = rng.r#gen());
         tensor.cast_to_dt(datum_type).unwrap().into_owned()
     }
 
diff --git a/src/tensor/mod.rs b/src/tensor/mod.rs
index 6efe007a..41d96e6f 100644
--- a/src/tensor/mod.rs
+++ b/src/tensor/mod.rs
@@ -62,7 +62,7 @@ pub trait TensorType: Clone + Debug + 'static {
 }
 
 macro_rules! tensor_type {
-    ($rust_type:ty, $tensor_type:ident, $zero:expr, $one:expr) => {
+    ($rust_type:ty, $tensor_type:ident, $zero:expr_2021, $one:expr_2021) => {
         impl TensorType for $rust_type {
             fn zero() -> Option<Self> {
                 Some($zero)
@@ -1693,7 +1693,7 @@ impl<T: TensorType + Rem<Output = T> + std::marker::Send + std::marker::Sync + P
         lhs.par_iter_mut()
             .zip(rhs)
             .map(|(o, r)| {
-                if let Some(zero) = T::zero() {
+                match T::zero() { Some(zero) => {
                     if r != zero {
                         *o = o.clone() % r;
                         Ok(())
@@ -1702,11 +1702,11 @@ impl<T: TensorType + Rem<Output = T> + std::marker::Send + std::marker::Sync + P
                             "Cannot divide by zero in remainder".to_string(),
                         ))
                     }
-                } else {
+                } _ => {
                     Err(TensorError::InvalidArgument(
                         "Undefined zero value".to_string(),
                     ))
-                }
+                }}
             })
             .collect::<Result<Vec<_>, _>>()?;
 
diff --git a/tests/py_integration_tests.rs b/tests/py_integration_tests.rs
index 353533b8..b9f5e13e 100644
--- a/tests/py_integration_tests.rs
+++ b/tests/py_integration_tests.rs
@@ -46,7 +46,9 @@ mod py_tests {
             assert!(status.success());
         });
         // set VOICE_DATA_DIR environment variable
-        std::env::set_var("VOICE_DATA_DIR", format!("{}", voice_data_dir));
+        unsafe {
+            std::env::set_var("VOICE_DATA_DIR", format!("{}", voice_data_dir));
+        }
     }
 
     fn download_catdog_data() {
@@ -63,7 +65,9 @@ mod py_tests {
             assert!(status.success());
         });
         // set VOICE_DATA_DIR environment variable
-        std::env::set_var("CATDOG_DATA_DIR", format!("{}", cat_and_dog_data_dir));
+        unsafe {
+            std::env::set_var("CATDOG_DATA_DIR", format!("{}", cat_and_dog_data_dir));
+        }
     }
 
     fn setup_py_env() {