From 6a0790ff5052d3d0cb652d2d9b50865b7f5e40e0 Mon Sep 17 00:00:00 2001 From: Starbeamrainbowlabs Date: Mon, 3 Oct 2022 16:32:09 +0100 Subject: [PATCH] convnext_inverse: add returns; change ordering --- aimodel/src/lib/ai/components/convnext_inverse.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/aimodel/src/lib/ai/components/convnext_inverse.py b/aimodel/src/lib/ai/components/convnext_inverse.py index 1fece03..8da69a9 100644 --- a/aimodel/src/lib/ai/components/convnext_inverse.py +++ b/aimodel/src/lib/ai/components/convnext_inverse.py @@ -29,6 +29,8 @@ def convnext_inverse(layer_in, depths, dims): for depth, dim in zip(depths, dims): layer_next = block_upscale(layer_next, i, depth=depth, dim=dim) i += 1 + + return layer_next def block_upscale(layer_in, block_number, depth, dim): @@ -36,5 +38,7 @@ def block_upscale(layer_in, block_number, depth, dim): for i in range(depth): layer_next = add_convnext_block(layer_next, dim=dim, prefix=f"cns.stage{block_number}.block.{i}") - layer_next = tf.keras.layers.LayerNormalization(name=f"cns.stage{block_number}.end.norm", epsilon=1e-6)(layer_next) layer_next = tf.keras.layers.Conv2DTranspose(name=f"cns.stage{block_number}.end.convtp", filters=dim, kernel_size=4, padding="same")(layer_next) + layer_next = tf.keras.layers.LayerNormalization(name=f"cns.stage{block_number}.end.norm", epsilon=1e-6)(layer_next) + + return layer_next