How to use the dnn.pytorch.layer function in dnn

def __init__(self, emb_matrix, args):
        """
        Initilize the model data and layer
        * emb_matrix [np.array]: word embedding matrix
        * args [dict]: all model arguments
        """
        nn.Module.__init__(self)
        base.base.__init__(self, args)

        self.emb_mat = layer.embedding_layer(emb_matrix, self.emb_type)
        self.pos_emb_mat = layer.positional_embedding_layer(self.n_hidden)
        self.drop_out = nn.Dropout(self.drop_prob)
        self.transformer = nn.ModuleList([
            layer.transformer_layer(self.emb_dim, self.n_hidden, self.n_head) for _ in range(self.n_layer)
        ])
        self.predict = layer.softmax_layer(self.emb_dim, self.n_class)

def init_weights(self):
        for m in self.modules():
            if isinstance(m, layer.transformer_layer):
                m.init_weights()
            if isinstance(m, layer.softmax_layer):
                m.init_weights()