Update to make test with cuda simpler.