Fix bug when using pissa method with deepspeed Former-commit-id: f260d458f91d6d2b4ed141f64844cded11d5aaad