[deepspeed] fix a bug in a test
This PR fixes a small bug in a test - was using zero3 for both stages by mistake, so fixing to test zero2 as well.

@sgugger 
_The documentation is not available anymore as the PR was closed or merged._