Coming from TensorFlow background, I am trying to convert a snippet of code of the custom layer from Keras to pytorch.
The custom layer in Keras looks like this:
class Attention_module(tf.keras.layers.Layer):
    def __init__(self, class_num):
        super(Attention_module self).__init__(class_num)
        self.class_num = class_num
        self.Ws = None

    def build(self, input_shape):
        embedding_length = int(input_shape[2])

        self.Ws = self.add_weight(shape=(self.class_num, embedding_length),
                                  initializer=tf.keras.initializers.get('glorot_uniform'), trainable=True)

        super(Attention_module, self).build(input_shape)

    def call(self, inputs):

        sentence_trans = tf.transpose(inputs, [0, 2, 1])
        at = tf.matmul(self.Ws, sentence_trans)
        at = tf.math.tanh(at)
        at = K.exp(at - K.max(at, axis=-1, keepdims=True))
        at = at / K.sum(at, axis=-1, keepdims=True)
        v = K.batch_dot(at, inputs)

        return v

I want to implement the same in the torch; I have already done the forward pass block but am confused about how to do the embedding and weight initialization the same as the above layer in pytorch?
class Attention_module(torch.nn.Module):
    def __init__(self, class_num):

        # how to initialize weight with same as above keras layer?

    def forward(self, inputs):
        
        sentence_trans = inputs.permute(0, 2, 1)

        at = torch.mm(self.Ws, sentence_trans)
        at = torch.nn.Tanh(at)
        at = torch.exp(at - torch.max(torch.Tensor(at), dim=-1, keepdims=True).values)
        at = at / torch.sum(at, dim = -1, keepdims=True)
        v = torch.einsum('ijk,ikl->ijl', at, inputs)

        return v

Thank you!