Apply suggestions from @sgugger's review

Commit

4 years ago

Apply suggestions from @sgugger's review Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>