Fix tests failing on a single GPU (#11753)

Commit

3 years ago

Fix tests failing on a single GPU (#11753) Co-authored-by: akihiro@grid.ai <akihiro@grid.ai> Co-authored-by: Carlos Mocholi <carlossmocholi@gmail.com>