[Feature Request] Addition of Kaiming He Normalization Option for Layer Initialization

[Zach-Attach]

🚀 Feature

Add a third option for weight_init as 'kaiming he' (see https://pytorch.org/docs/stable/nn.init.html and https://arxiv.org/pdf/1803.08494).

Motivation

I have been using this method of initialization for a while as it performs better than default or orthogonal. I thought it would make sens to add this to the public repo.sense

Pitch

Already wrote the code changes here.
5b08294

Alternatives

No response

Additional context

Tests do not currently exist for layer initialization, so I haven't added any for such a simple addition.

Checklist

• I have checked that there is no similar issue in the repo