diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 8ef5d447d4..4bc7098cf8 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -107,41 +107,6 @@ jobs:
           key: ${{ runner.os }}-${{ matrix.cache }}-${{ matrix.test}}-${{ hashFiles('**/Cargo.toml') }}
           prefix-key: "v5-rust"
 
-      - name: free disk space
-        if: runner.os == 'Linux'
-        run: |
-          df -h
-          sudo swapoff -a
-          sudo rm -f /swapfile
-          sudo apt clean
-          df -h
-          cargo clean --package burn-tch
-
-      - name: install llvmpipe and lavapipe
-        if: runner.os == 'Linux'
-        uses: ./.github/actions/setup-llvmpipe-lavapipe
-
-      - name: Run cargo clippy for stable version
-        if: runner.os == 'Linux' && matrix.rust == 'stable' && matrix.test == 'std'
-        uses: giraffate/clippy-action@v1
-        with:
-          github_token: ${{ secrets.GITHUB_TOKEN }}
-          # Run clippy for each workspace, targets, and featrues, considering
-          # warnings as errors
-          clippy_flags: --all-targets -- -Dwarnings
-          # Do not filter results
-          filter_mode: nofilter
-          # Report clippy annotations as snippets
-          reporter: github-pr-check
-
-      - name: Install grcov
-        if: runner.os == 'Linux' && matrix.rust == 'stable' && matrix.test == 'std'
-        env:
-          GRCOV_LINK: https://github.com/mozilla/grcov/releases/download
-        run: |
-          curl -L "$GRCOV_LINK/v$GRCOV_VERSION/grcov-x86_64-unknown-linux-musl.tar.bz2" |
-          tar xj -C $HOME/.cargo/bin
-
       # -----------------------------------------------------------------------------------
       # BEGIN -- Windows steps disabled as long as DISABLE_WGPU=1 (wgpu tests are disabled)
       # -----------------------------------------------------------------------------------
@@ -198,57 +163,6 @@ jobs:
       # END -- Windows steps disabled as long as DISABLE_WGPU=1 (wgpu tests are disabled)
       # -----------------------------------------------------------------------------------
 
-      - name: (linux) install vulkan sdk
-        # from wgpu repo: https://github.com/gfx-rs/wgpu/blob/trunk/.github/workflows/ci.yml
-        if: runner.os == 'Linux'
-        shell: bash
-        run: |
-          set -e
-
-          sudo apt-get update -y -qq
-
-          # vulkan sdk
-          wget -qO - https://packages.lunarg.com/lunarg-signing-key-pub.asc | sudo apt-key add -
-          sudo wget -qO /etc/apt/sources.list.d/lunarg-vulkan-$VULKAN_SDK_VERSION-jammy.list https://packages.lunarg.com/vulkan/$VULKAN_SDK_VERSION/lunarg-vulkan-$VULKAN_SDK_VERSION-jammy.list
-
-          sudo apt-get update
-          sudo apt install -y vulkan-sdk
-
-      - name: (linux) install mesa
-        # from wgpu repo: https://github.com/gfx-rs/wgpu/blob/trunk/.github/workflows/ci.yml
-        if: runner.os == 'Linux'
-        shell: bash
-        run: |
-          set -e
-
-          curl -L --retry 5 https://github.com/gfx-rs/ci-build/releases/download/$CI_BINARY_BUILD/mesa-$MESA_VERSION-linux-x86_64.tar.xz -o mesa.tar.xz
-          mkdir mesa
-          tar xpf mesa.tar.xz -C mesa
-
-          # The ICD provided by the mesa build is hardcoded to the build environment.
-          #
-          # We write out our own ICD file to point to the mesa vulkan
-          cat <<- EOF > icd.json
-          {
-            "ICD": {
-                "api_version": "1.1.255",
-                "library_path": "$PWD/mesa/lib/x86_64-linux-gnu/libvulkan_lvp.so"
-            },
-            "file_format_version": "1.0.0"
-          }
-          EOF
-
-          echo "VK_DRIVER_FILES=$PWD/icd.json" >> "$GITHUB_ENV"
-          echo "LD_LIBRARY_PATH=$PWD/mesa/lib/x86_64-linux-gnu/:$LD_LIBRARY_PATH" >> "$GITHUB_ENV"
-          echo "LIBGL_DRIVERS_PATH=$PWD/mesa/lib/x86_64-linux-gnu/dri" >> "$GITHUB_ENV"
-
       - name: run checks & tests
         shell: bash
-        run: ${{ matrix.coverage-flags }} ${{ matrix.wgpu-flags }} cargo xtask run-checks ${{ matrix.test }}
-
-      - name: Codecov upload
-        if: runner.os == 'Linux' && matrix.rust == 'stable' && matrix.test == 'std'
-        uses: codecov/codecov-action@v4
-        with:
-          files: lcov.info
-          token: ${{ secrets.CODECOV_TOKEN }}
+        run: DISABLE_WGPU=1 cargo xtask run-checks std
diff --git a/crates/burn-core/src/nn/attention/mask.rs b/crates/burn-core/src/nn/attention/mask.rs
index b1c277f0ee..73373bc8ff 100644
--- a/crates/burn-core/src/nn/attention/mask.rs
+++ b/crates/burn-core/src/nn/attention/mask.rs
@@ -88,11 +88,12 @@ pub fn generate_padding_mask<B: Backend>(
 
 #[cfg(test)]
 mod tests {
-    use super::*;
-    use crate::TestBackend;
-    use alloc::vec;
     use burn_tensor::Data;
 
+    use crate::TestBackend;
+
+    use super::*;
+
     #[test]
     fn test_generate_autoregressive_mask() {
         let device = <TestBackend as Backend>::Device::default();