Flux save and load model

I’m saving the UNet model as
function save_unet_train_state(filename::String, model::Unet, opt, loss::Function) @assert endswith(filename, ".jld2") @save filename model opt loss
where opt is opt = RMSProp(learning_rate, momentum)
but on loading, i wrote `

> model_loaded, opt_loaded, loss_loaded = UNet.load_unet(joinpath(pwd(), "UNet.jld2"))
> ┌ Warning: saved type Unet is missing field labels in workspace type; reconstructing
> └ @ JLD2 ~/.julia/packages/JLD2/k9Gt0/src/data/reconstructing_datatypes.jl:152
> ERROR: MethodError: no method matching iterate(::JLD2.ReconstructedTypes.var"##Unet#298")
> Closest candidates are:
>   iterate(::Union{LinRange, StepRangeLen}) at range.jl:664
>   iterate(::Union{LinRange, StepRangeLen}, ::Int64) at range.jl:664

Not sure how to fix this?
Also, the field label is a vector of classes, e.g. [1,2,3]. I thought saving model will save it as model.labels automatically?

You might have to run all of the using and includes that were present when you saved the module, so that the deserialiser knows how to construct the objects.

EDIT: I am not familiar with UNet, but shouldn’t you be using a @load to load the data back? Similar to how it is saved in the first place?

Sorry I didn’t include the original definition of load.
@load filename model opt loss labels
However, the @save didn’t save one field which is labels.
Labels was added at training when model was initialized
model = Unet(num_channels, labels) |> device