Wrap the layers in the VGG type

FluxML · Jan 4, 2024 · 1207f84 · 1207f84
1 parent 09f91c0
commit 1207f84
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/src/convnets/vgg.jl b/src/convnets/vgg.jl
@@ -130,7 +130,8 @@ end
 function VGG(depth::Integer; pretrain::Bool = false, batchnorm::Bool = false,
              inchannels::Integer = 3, nclasses::Integer = 1000)
     _checkconfig(depth, keys(VGG_CONFIGS))
-    model = vgg((224, 224); config = VGG_CONFIGS[depth], batchnorm, inchannels, nclasses)
+    layers = vgg((224, 224); config = VGG_CONFIGS[depth], batchnorm, inchannels, nclasses)
+    model = VGG(layers)
     if pretrain
         artifact_name = string("vgg", depth)
         if batchnorm