There are several aspects:

the parameter is created by assigning self.Ws as a  nn.Parameter of the appropriate shape in __init__ (after super().__init__()).
the equivalent initialization is torch.nn.init.xavier_uniform_, but you may have to adjust the gain parameter. Check the mean and std (and a histogram plot)
you have some funny stuff in the forward mixing torch.nn things where I would have expected torch.nn.functional, also you should never need to use  torch.Tensor,
finally, are you spelling out torch.nn.functional.softmax in the lines with torch.exp and the following? I’d probably write it as softmax if you are. I’m not entirely sure why you would to tanh and then softmax, but hey, what do I know of your problem.

I didn’t look at the dimension ordering (the permute and others).
Best regards
Thomas