diff --git a/model_tf2.py b/model_tf2.py
index 88793ed..eae8f66 100644
--- a/model_tf2.py
+++ b/model_tf2.py
@@ -375,7 +375,7 @@ def Deeplabv3(weights='pascal_voc', input_tensor=None, input_shape=(512, 512, 3)
     # upsample. have to use compat because of the option align_corners
     size_before = K.int_shape(x)
 
-    b4 = Lambda(lambda x: tf.image.resize_with_pad(x, size_before[1:3], method='bilinear'))(b4)
+    b4 = Lambda(lambda x: tf.image.resize_with_pad(x, target_height=size_before[1], target_width=size_before[2]))(b4)
     # b4 = UpSampling2D(size=(size_before[1],size_before[2]),interpolation='bilinear')(b4)
     # simple 1x1
     b0 = Conv2D(256, (1, 1), padding='same', use_bias=False, name='aspp0')(x)
@@ -412,7 +412,7 @@ def Deeplabv3(weights='pascal_voc', input_tensor=None, input_shape=(512, 512, 3)
         # size_in = K.int_shape(x)
         # size_out = K.int_shape(skip1)
         # x = UpSampling2D(size=(size_out[1]//size_in[1],size_out[2]//size_in[2]),interpolation='bilinear')(x)
-        x = Lambda(lambda xx: tf.image.resize_with_pad(xx, skip1.shape[1:3], method='bilinear'))(x)
+        x = Lambda(lambda xx: tf.image.resize_with_pad(xx, target_height=skip1.shape[1], target_width=skip1.shape[2]))(x)
         dec_skip1 = Conv2D(48, (1, 1), padding='same',
                            use_bias=False, name='feature_projection0')(skip1)
         dec_skip1 = BatchNormalization(
@@ -436,7 +436,7 @@ def Deeplabv3(weights='pascal_voc', input_tensor=None, input_shape=(512, 512, 3)
     # size_out = K.int_shape(img_input)
     # x = UpSampling2D(size=(size_out[1] // size_in[1], size_out[2] // size_in[2]), interpolation='bilinear')(x)
     size_before3 = K.int_shape(img_input)
-    x = Lambda(lambda xx: tf.image.resize_with_pad(xx, size_before3[1:3], method='bilinear'))(x)
+    x = Lambda(lambda xx: tf.image.resize_with_pad(xx, target_height=size_before3[1], target_width=size_before3[2]))(x)
     # Ensure that the model takes into account
     # any potential predecessors of `input_tensor`.
     if input_tensor is not None: