ROCm Support PR Proposal

[reimorster]

I proposed small changes on main script to allow people with ROCm hardware to iterate locally or on cheaper cloud runpods, even though the final submission must attend challenge rules and run on 8xH100 cluster.
It would be nice if someone reviews it.

#735

PS. @0hq , since theres a CUDA and MLX support, it would be fair to also bring AMD to the game.

Thanks