OUTPUT_CHANNELS = 3
base_model = tf.keras.applications.MobileNetV2(input_shape=[256, 256, 3], include_top=False)

Use the activations of these layers
layer_names = [
   'block_1_expand_relu',   # 64x64
   'block_3_expand_relu',   # 32x32
   'block_6_expand_relu',   # 16x16
   'block_13_expand_relu',  # 8x8
   'block_16_project',      # 4x4
#]
#layers = [base_model.get_layer(name).output for name in layer_names]

# Create the feature extraction model
#down_stack = tf.keras.Model(inputs=base_model.input, outputs=layers)

#down_stack.trainable = False