update ENV to ROCm dockers

[HaiShaw]

Motivation

Performance boost

Modifications

As it is.

Checklist

• [+] Format your code according to the Code Formatting with Pre-Commit.
• [+] Add unit tests as outlined in the Running Unit Tests.
• [+] Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• [+] Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling.