fix: try disabling force_preserve

LuxDL · Sep 13, 2024 · c30b4da · c30b4da
1 parent 4f1d863
commit c30b4da
Showing 1 changed file with 8 additions and 8 deletions.
diff --git a/ext/BoltzMetalheadExt.jl b/ext/BoltzMetalheadExt.jl
@@ -11,32 +11,32 @@ Utils.is_extension_loaded(::Val{:Metalhead}) = true
 
 function Vision.ResNetMetalhead(depth::Int; pretrained::Bool=false)
     @argcheck depth in (18, 34, 50, 101, 152)
-    return FromFluxAdaptor(; preserve_ps_st=pretrained, force_preserve=true)(Metalhead.ResNet(
+    return FromFluxAdaptor(; preserve_ps_st=pretrained)(Metalhead.ResNet(
         depth; pretrain=pretrained).layers)
 end
 
 function Vision.ResNeXtMetalhead(
         depth::Int; cardinality=32, base_width=nothing, pretrained::Bool=false)
     @argcheck depth in (50, 101, 152)
     base_width = base_width === nothing ? (depth == 101 ? 8 : 4) : base_width
-    return FromFluxAdaptor(; preserve_ps_st=pretrained, force_preserve=true)(Metalhead.ResNeXt(
+    return FromFluxAdaptor(; preserve_ps_st=pretrained)(Metalhead.ResNeXt(
         depth; pretrain=pretrained, cardinality, base_width).layers)
 end
 
 function Vision.GoogLeNetMetalhead(; pretrained::Bool=false)
-    return FromFluxAdaptor(; preserve_ps_st=pretrained, force_preserve=true)(Metalhead.GoogLeNet(;
+    return FromFluxAdaptor(; preserve_ps_st=pretrained)(Metalhead.GoogLeNet(;
         pretrain=pretrained).layers)
 end
 
 function Vision.DenseNetMetalhead(depth::Int; pretrained::Bool=false)
     @argcheck depth in (121, 161, 169, 201)
-    return FromFluxAdaptor(; preserve_ps_st=pretrained, force_preserve=true)(Metalhead.DenseNet(
+    return FromFluxAdaptor(; preserve_ps_st=pretrained)(Metalhead.DenseNet(
         depth; pretrain=pretrained).layers)
 end
 
 function Vision.MobileNetMetalhead(name::Symbol; pretrained::Bool=false)
     @argcheck name in (:v1, :v2, :v3_small, :v3_large)
-    adaptor = FromFluxAdaptor(; preserve_ps_st=pretrained, force_preserve=true)
+    adaptor = FromFluxAdaptor(; preserve_ps_st=pretrained)
     model = if name == :v1
         adaptor(Metalhead.MobileNetv1(; pretrain=pretrained).layers)
     elseif name == :v2
@@ -51,18 +51,18 @@ end
 
 function Vision.ConvMixerMetalhead(name::Symbol; pretrained::Bool=false)
     @argcheck name in (:base, :large, :small)
-    return FromFluxAdaptor(; preserve_ps_st=pretrained, force_preserve=true)(Metalhead.ConvMixer(
+    return FromFluxAdaptor(; preserve_ps_st=pretrained)(Metalhead.ConvMixer(
         name; pretrain=pretrained).layers)
 end
 
 function Vision.SqueezeNetMetalhead(; pretrained::Bool=false)
-    return FromFluxAdaptor(; preserve_ps_st=pretrained, force_preserve=true)(Metalhead.SqueezeNet(;
+    return FromFluxAdaptor(; preserve_ps_st=pretrained)(Metalhead.SqueezeNet(;
         pretrain=pretrained).layers)
 end
 
 function Vision.WideResNetMetalhead(depth::Int; pretrained::Bool=false)
     @argcheck depth in (18, 34, 50, 101, 152)
-    return FromFluxAdaptor(; preserve_ps_st=pretrained, force_preserve=true)(Metalhead.WideResNet(
+    return FromFluxAdaptor(; preserve_ps_st=pretrained)(Metalhead.WideResNet(
         depth; pretrain=pretrained).layers)
 end