tests/nn/data_parallel/test_fsdp.py · fa1b85fbbe75f058b39f1bcf027de42e6ddbd487 · OpenDAS / fairscale

[fix][FSDP] fix weight init when using apply() (fixes #490 and #444) (#543) · fa1b85fb

Myle Ott authored Mar 20, 2021

* Add new test for weight init (fails)
* Set FSDP.compute_device so summon_full_params works before module moves to CUDA
* Override FSDP.apply to enable custom weight init

fa1b85fb

test_fsdp.py 37.3 KB

Replace test_fsdp.py