Ipanfilo/ci test fixes #775

Job	Run time
determine_level	4s
CI Level 3 / Select Docker Image	5s
CI Level 3 / build / Build ROCm Docker image and TransformerEngine wheels	23m 27s
CI Level 3 / mGPU Torch (mi30x)	54m 13s
CI Level 3 / sGPU Tests (mi35x)	5h 33m 8s
CI Level 3 / mGPU JAX (mi35x)	35m 29s
CI Level 3 / sGPU Tests (mi30x)	2h 53m 33s
CI Level 3 / mGPU JAX (mi30x)	51m 25s
CI Level 3 / mGPU Torch (mi35x)	35m 5s
determine_level	4s
CI Level 3 / sGPU Tests (mi35x)	2h 25m 9s
CI Level 3 / mGPU JAX (mi35x)	27m 39s
CI Level 3 / Select Docker Image	5s
CI Level 3 / sGPU Tests (mi30x)	2h 53m 33s
CI Level 3 / build / Build ROCm Docker image and TransformerEngine wheels	23m 27s
CI Level 3 / mGPU Torch (mi30x)	54m 13s
CI Level 3 / mGPU JAX (mi30x)	51m 25s
CI Level 3 / mGPU Torch (mi35x)	35m 5s
	20h 17m 9s

Provide feedback