Flux.loadmodel! does not work with old neural_snow model states

[imreddyTeja]

Describe the bug

The get_znetwork() and tests/standalone/Snow/tool_tests.jl error when Flux.loadmodel! is called.

It also seems like this test is not exercised with CI.

The test also looks like it was not updated after recent PRs that changed the snow model.

There is also a deprecation of implicitly taking the gradient with Flux.jl. This requires some changed to the way the model is setup.

Lastly, NeuralSnowExt is conditionally loaded with cuDNN. cuDNN is no longer maintained, and its functionality is replaced by cuArrays, which is deprecated and replaced by CUDA.

To Reproduce

Run tests/standalone/Snow/tool_tests.jl with CUDA

julia> include("test/standalone/Snow/tool_tests.jl")
Test Summary:          | Pass  Total  Time
Testing Data Utilities |   43     43  1.6s
Testing Model Utilities: Error During Test at /home/treddy/ClimaLand.jl/test/standalone/Snow/tool_tests.jl:213
  Got exception outside of a @test
  ArgumentError: Tried to load (:layers,) into (:layers, :connection) but the structures do not match.

Or even better, add it as a unit test, and open pull request.

Project

If not using the `examples` project: ``` paste your Project.toml here. ``` ``` paste your Manifest.toml here. ```

System details

Ran on Clima cluster.

Related issues / PRs

relates to #974

Similar issues in other repos: issue

Please add any relevant links.