tree 79839f19defe69c3140d0b64289524e397e2df8e
parent cedfa4ccd89d3b2869e8433cc8a02ba6afa8e8c2
author Pritam Damania <pritam.damania@fb.com> 1611864472 -0800
committer Facebook GitHub Bot <facebook-github-bot@users.noreply.github.com> 1611864621 -0800

[Pipe] Refactor convert_to_balance under non-test package. (#50860)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/50860

Since fairscale.nn.Pipe still uses 'balance' and 'devices' parameters,
other frameworks like fairseq still use these parameters. As a result, the
`convert_to_balance` method is a nice utility to use for migrating to PyTorch
Pipe without changing a lot of code in other frameworks.

In addition to this I've renamed the method to be more illustrative of what it
does and also allowed an optional devices parameter.
ghstack-source-id: 120430775

Test Plan:
1) waitforbuildbot
2) Tested with fairseq

Reviewed By: SciPioneer

Differential Revision: D25987273

fbshipit-source-id: dccd42cf1a74b08c876090d3a10a94911cc46dd8