Okay one last Question how to write the below function in Pytorch’s way
def GlobalAveragePooling2D_r(f):
    def func(x):
    
        repc =  int(x.shape[4])
        m    =  keras.backend.repeat_elements(f, repc, axis = 4)
        x    =  layers.multiply([x, m])
        repx =  int(x.shape[2])
        repy =  int(x.shape[3])
        x    = (keras.backend.sum(x, axis=[2, 3], keepdims=True) / (keras.backend.sum(m, axis=[2, 3], keepdims=True)))
        x    =  keras.backend.repeat_elements(x, repx, axis = 2)
        x    =  keras.backend.repeat_elements(x, repy, axis = 3)    
        return x
    return Lambda(func)