test: more recurrent testing fixes

LuxDL · Sep 22, 2024 · c27fcfc · c27fcfc
1 parent f135943
commit c27fcfc
Showing 1 changed file with 7 additions and 7 deletions.
diff --git a/test/layers/recurrent_tests.jl b/test/layers/recurrent_tests.jl
@@ -2,8 +2,8 @@
 
 using MLDataDevices
 
-MLDataDevices.get_device_type(::Function) = Nothing # FIXME: upstream maybe?
-MLDataDevices.get_device_type(_) = Nothing # FIXME: upstream maybe?
+MLDataDevices.Internal.get_device_type(::Function) = Nothing # FIXME: upstream maybe?
+MLDataDevices.Internal.get_device_type(_) = Nothing # FIXME: upstream maybe?
 
 function loss_loop(cell, x, p, st)
  (y, carry), st_ = cell(x, p, st)
@@ -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