ROCm Support PR Proposal

I proposed small changes on main script to allow people with ROCm hardware to iterate locally even though the final result must run on H100 cluster. It would be nice if someone reviews it.

Thanks