diff --git a/README.md b/README.md
index 3f899126cf..396839b7c3 100644
--- a/README.md
+++ b/README.md
@@ -11,8 +11,9 @@
 [![Stable Docs](https://img.shields.io/badge/docs-stable-blue.svg)](http://lux.csail.mit.edu/stable/)
 
 [![CI](https://github.com/LuxDL/Lux.jl/actions/workflows/CI.yml/badge.svg?branch=main)](https://github.com/LuxDL/Lux.jl/actions/workflows/CI.yml)
-[![CI (pre-release)](https://github.com/LuxDL/Lux.jl/actions/workflows/CIPreRelease.yml/badge.svg?branch=main)](https://github.com/LuxDL/Lux.jl/actions/workflows/CIPreRelease.yml)
-[![Build status](https://img.shields.io/buildkite/ba1f9622add5978c2d7b194563fd9327113c9c21e5734be20e/main.svg?label=gpu&branch=main)](https://buildkite.com/julialang/lux-dot-jl)
+[![CI (pre-release)](https://img.shields.io/github/actions/workflow/status/LuxDL/Lux.jl/CIPreRelease.yml?branch=main&label=CI%20(pre-release)&logo=github)](https://github.com/LuxDL/Lux.jl/actions/workflows/CIPreRelease.yml)
+[![Build status](https://img.shields.io/buildkite/ba1f9622add5978c2d7b194563fd9327113c9c21e5734be20e/main.svg?label=gpu&branch=main&logo=buildkite)](https://buildkite.com/julialang/lux-dot-jl)
+
 [![codecov](https://codecov.io/gh/LuxDL/Lux.jl/branch/main/graph/badge.svg?token=IMqBM1e3hz)](https://codecov.io/gh/LuxDL/Lux.jl)
 [![Benchmarks](https://github.com/LuxDL/Lux.jl/actions/workflows/Benchmark.yml/badge.svg?branch=main)](https://lux.csail.mit.edu/benchmarks/)
 
diff --git a/test/layers/basic_tests.jl b/test/layers/basic_tests.jl
index a09cd1b2ef..8fcca38005 100644
--- a/test/layers/basic_tests.jl
+++ b/test/layers/basic_tests.jl
@@ -21,9 +21,9 @@
             layer2 = ReverseSequence(2)
             layer3 = ReverseSequence(1)
             display(layer)
-            ps, st = Lux.setup(rng, layer) |> device
-            ps2, st2 = Lux.setup(rng, layer2) |> device
-            ps3, st3 = Lux.setup(rng, layer3) |> device
+            ps, st = Lux.setup(rng, layer) .|> device
+            ps2, st2 = Lux.setup(rng, layer2) .|> device
+            ps3, st3 = Lux.setup(rng, layer3) .|> device
 
             x = randn(rng, 3) |> aType
             xr = reverse(x)
@@ -440,44 +440,27 @@ end
             @jet layer(x, ps, st)
 
             x = (rand(1:vocab_size[1], 3), rand(1:vocab_size[2], 3)) .|> aType
-            if mode == "cuda"
-                @test_broken begin
-                    y, st_ = layer(x, ps, st)
-                    @test y isa aType{Float32}
-                    @test y == ps.weight[:, CartesianIndex.(x...)]
-                end
-            else
-                y, st_ = layer(x, ps, st)
-                @test y isa aType{Float32}
-                @test y == ps.weight[:, CartesianIndex.(x...)]
-            end
+            y, st_ = layer(x, ps, st)
+            @test y isa aType{Float32}
+            @test y == ps.weight[:, CartesianIndex.(x...)]
+
             @jet layer(x, ps, st)
 
-            if mode == "cuda"
-                @test_broken begin
-                    x = (rand(1:vocab_size[1], 3, 4), rand(1:vocab_size[2], 3, 4)) .|> aType
-                    y, st_ = layer(x, ps, st)
-                    @test y isa aType{Float32, 3}
-                    @test size(y) == (embed_size, 3, 4)
-                end
-            else
-                x = (rand(1:vocab_size[1], 3, 4), rand(1:vocab_size[2], 3, 4)) .|> aType
-                y, st_ = layer(x, ps, st)
-                @test y isa aType{Float32, 3}
-                @test size(y) == (embed_size, 3, 4)
-            end
+            x = (rand(1:vocab_size[1], 3, 4), rand(1:vocab_size[2], 3, 4)) .|> aType
+            y, st_ = layer(x, ps, st)
+            @test y isa aType{Float32, 3}
+            @test size(y) == (embed_size, 3, 4)
+
             @jet layer(x, ps, st)
 
-            if mode != "cuda"
-                x = (rand(1:vocab_size[1], 3), rand(1:vocab_size[2], 4)) .|> aType
-                @test_throws DimensionMismatch layer(x, ps, st)
+            x = (rand(1:vocab_size[1], 3), rand(1:vocab_size[2], 4)) .|> aType
+            @test_throws DimensionMismatch layer(x, ps, st)
 
-                x = (rand(1:vocab_size[1], 3, 4), rand(1:vocab_size[2], 4, 5)) .|> aType
-                @test_throws DimensionMismatch layer(x, ps, st)
+            x = (rand(1:vocab_size[1], 3, 4), rand(1:vocab_size[2], 4, 5)) .|> aType
+            @test_throws DimensionMismatch layer(x, ps, st)
 
-                x = ()
-                @test_throws ArgumentError layer(x, ps, st)
-            end
+            x = ()
+            @test_throws ArgumentError layer(x, ps, st)
         end
     end
 end