transformer_engine/pytorch/module/linear.py · c32a62ccb0a38d14d98e8bd6a14ace2d73bbac25 · OpenDAS / TransformerEngine

Add CPU initialization method (#368) · 2451c8d1

Kirthi Shankar Sivamani authored Aug 16, 2023



* CPU initialization
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Fix default value
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

* Change API and add to RMSNorm
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

---------
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>

2451c8d1

linear.py 29.4 KB

Replace linear.py