diff --git a/test/test_amdgpu.jl b/test/test_amdgpu.jl
index 86e86ff4b6..cc63063b9f 100644
--- a/test/test_amdgpu.jl
+++ b/test/test_amdgpu.jl
@@ -1,25 +1,36 @@
 include("dependencies_for_runtests.jl")
 
 using AMDGPU
+using SeawaterPolynomials.TEOS10: TEOS10EquationOfState
 
-@testset "AMDGPU extension" begin
+@testset "AMDGPU with HydrostaticFreeSurfaceModel" begin
     roc = AMDGPU.ROCBackend()
     arch = GPU(roc)
 
     for FT in float_types
         @info "    Testing on $arch with $FT"
 
-        grid = RectilinearGrid(arch, FT, size=(4, 8, 16), x=[0, 1, 2, 3, 4], y=(0, 1), z=(0, 16))
+        grid = LatitudeLongitudeGrid(arch, FT, size=(4, 8, 16), longitude=(-60, 60), latitude=(0, 60), z=(0, 1))
 
-        @test parent(grid.xᶠᵃᵃ) isa ROCArray
-        @test parent(grid.xᶜᵃᵃ) isa ROCArray
+        @test parent(grid.Δxᶜᶜᵃ) isa ROCArray
+        @test parent(grid.Δxᶠᶜᵃ) isa ROCArray
+        @test parent(grid.Δxᶜᶠᵃ) isa ROCArray
+        @test parent(grid.Δxᶠᶠᵃ) isa ROCArray
+        @test parent(grid.Δyᶠᶜᵃ) isa ROCArray
+        @test parent(grid.Δyᶜᶠᵃ) isa ROCArray
+        @test parent(grid.Azᶜᶜᵃ) isa ROCArray
+        @test parent(grid.Azᶠᶜᵃ) isa ROCArray
+        @test parent(grid.Azᶜᶠᵃ) isa ROCArray
+        @test parent(grid.Azᶠᶠᵃ) isa ROCArray
         @test eltype(grid) == FT
         @test architecture(grid) isa GPU
 
-        model = HydrostaticFreeSurfaceModel(; grid,
+        equation_of_state = TEOS10EquationOfState()
+        buoyancy = SeawaterBuoyancy(; equation_of_state)
+
+        model = HydrostaticFreeSurfaceModel(; grid, buoyancy,
                                             coriolis = FPlane(latitude=45),
-                                            buoyancy = BuoyancyTracer(),
-                                            tracers = :b,
+                                            tracers = (:T, :S),
                                             momentum_advection = WENO(order=5),
                                             tracer_advection = WENO(order=5),
                                             free_surface = SplitExplicitFreeSurface(grid; substeps=60))
@@ -35,3 +46,36 @@ using AMDGPU
         @test time(simulation) == 3minutes
     end
 end
+
+@testset "AMDGPU with NonhydrostaticModel" begin
+    roc = AMDGPU.ROCBackend()
+    arch = GPU(roc)
+
+    for FT in float_types
+        @info "    Testing on $arch with $FT"
+
+        z = 0:16 |> collect
+        grid = RectilinearGrid(arch, FT, size=(4, 8, 16), x=(0, 1), y=(0, 1), z=z)
+
+        @test parent(grid.z.cᵃᵃᶠ) isa ROCArray
+        @test parent(grid.z.cᵃᵃᶜ) isa ROCArray
+        @test eltype(grid) == FT
+        @test architecture(grid) isa GPU
+
+        model = NonhydrostaticModel(; grid,
+                                    coriolis = FPlane(latitude=45),
+                                    buoyancy = BuoyancyTracer(),
+                                    tracers = :b,
+                                    advection = WENO(order=5))
+
+        for field in merge(model.velocities, model.tracers)
+            @test parent(field) isa ROCArray
+        end
+
+        simulation = Simulation(model, Δt=1minute, stop_iteration=3)
+        run!(simulation)
+
+        @test iteration(simulation) == 3
+        @test time(simulation) == 3minutes
+    end
+end