Update models/unet_gray2color.py

models/unet_gray2color.py

@@ -2,9 +2,9 @@ import numpy as np
 import tensorflow as tf
 from tensorflow.keras.layers import (
     Input, Dense, Conv2D, MaxPooling2D, UpSampling2D, Concatenate,
-    BatchNormalization, LayerNormalization, Dropout, MultiHeadAttention, Add, Reshape
+    BatchNormalization, LayerNormalization, Dropout, MultiHeadAttention, Add, Reshape, Layer
 )
-from tensorflow.keras.models import Model
+from tensorflow.keras.models import Model, layers
 from tensorflow.keras.optimizers import Adam
 from tensorflow.keras.callbacks import ModelCheckpoint, ReduceLROnPlateau, EarlyStopping
 from tensorflow.keras.mixed_precision import set_global_policy