Weight initialization with specific distribution

I want to initialize the weights of my neural network with parameters samples from a specific distribution, that is not already present in nn.init module. Let’s say for example a beta distribution.

I saw that you can use the method described there: python - Custom weight initialization in PyTorch - Stack Overflow

But there you need to use the nn.init. What is an elegant way to initialize weights from a distribution that I specify?

You can customize it by referring to the initialization source code here: pytorch/init.py at master · pytorch/pytorch · GitHub.