[cuda ext tests] fixing tests
Fixing several tests for cuda extension jobs single and multi-gpu.

One failure I couldn't reproduce locally, posted about it here https://github.com/huggingface/transformers/issues/11618

@LysandreJik or @sgugger  