Mildnet for pytorch?

Hello, I am new in Pytorch and would need a little bit of help.
I am trying to use Mildnet architecture for image retrieval task. I have the network defined in
Tensorflow estimator api and would like to use it as a pytoch model.
Paper is at : https://arxiv.org/abs/1903.00905
Code for Keras is at : https://github.com/gofynd/mildnet

My code is for Tensorflow Keras API

module = hub.Module(https://tfhub.dev/google/imagenet/mobilenet_v1_100_224/feature_vector/1, trainable=True)
out = module(images, signature="image_feature_vector", as_dict=True)


out1 = out['MobilenetV1/Conv2d_1_depthwise']
out2 = out['MobilenetV1/Conv2d_2_depthwise']
out4 = out['MobilenetV1/Conv2d_4_depthwise']
out6 = out['MobilenetV1/Conv2d_6_depthwise']
out12 = out['MobilenetV1/Conv2d_12_depthwise']
out13 = out['MobilenetV1/Conv2d_13_depthwise']

intermediate_layer_outputs = [out1, out2, out4, out6, out12]
convnet_output = layers.GlobalAveragePooling2D()(out13)
for output in intermediate_layer_outputs:
  output = layers.GlobalAveragePooling2D()(output)
  convnet_output = layers.concatenate([convnet_output, output])

convnet_output = layers.Dense(1024, activation='relu')(convnet_output)
convnet_output = layers.Dropout(0.5)(convnet_output)
convnet_output = layers.Dense(128, activation='relu')(convnet_output)
#embeddings = layers.Lambda(lambda  x: K.l2_normalize(x,axis=1))(convnet_output)
if params.l2_normalize:
    convnet_output = tf.nn.l2_normalize(convnet_output, axis=1)

return convnet_output

Would you help me to make it work? Thanks a lot