tests/test_modeling_megatron_bert.py · d4e4efce68a5d18ca3175475b59051dec336fa6b · chenpangpang / transformers

Initial support for symbolic tracing with torch.fx allowing dynamic axes (#13579) · d4e4efce

Michael Benayoun authored Oct 05, 2021



* Symbolic trace dynamic axes support for BERT like models (albert, bert, distilbert, mobilebert, electra, megatron-bert)
* Sanity checks before tracing that make sure the model to trace is supported
* Adapted to PyTorch 1.9
Co-authored-by: Michael Benayoun <michael@huggingface.co>

d4e4efce

test_modeling_megatron_bert.py 15.6 KB

Replace test_modeling_megatron_bert.py