test: more recurrent testing fixes

LuxDL · Sep 22, 2024 · d8bbfc2 · d8bbfc2
1 parent f135943
commit d8bbfc2
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/test/layers/recurrent_tests.jl b/test/layers/recurrent_tests.jl
@@ -43,9 +43,9 @@ end
  @jet rnncell((x, carry), ps, st)
 
  if train_state
- @test hasproperty(ps, :train_state)
+ @test hasproperty(ps, :hidden_state)
  else
- @test !hasproperty(ps, :train_state)
+ @test !hasproperty(ps, :hidden_state)
  end
 
  @test_gradients(loss_loop, rnncell, x, ps, st; atol=1.0f-3, rtol=1.0f-3)
@@ -95,8 +95,8 @@ end
  @jet lstmcell(x, ps, st)
  @jet lstmcell((x, carry), ps, st)
 
- @test !hasproperty(ps, :train_state)
- @test !hasproperty(ps, :train_memory)
+ @test !hasproperty(ps, :hidden_state)
+ @test !hasproperty(ps, :memory)
 
  @test_gradients(loss_loop, lstmcell, x, ps, st; atol=1.0f-3, rtol=1.0f-3)
  end
@@ -198,7 +198,7 @@ end
  @jet grucell(x, ps, st)
  @jet grucell((x, carry), ps, st)
 
- @test !hasproperty(ps, :train_state)
+ @test !hasproperty(ps, :hidden_state)
 
  @test_gradients(loss_loop, grucell, x, ps, st; atol=1e-3, rtol=1e-3)
  end