[ROCM] Fix xla_call_module_test
Hi @cheshire Can you please review this PR ? Thank you! @i-chaochen can you address merge conflicts? > @i-chaochen can you address merge conflicts?

Thanks @reedwm for the notice, it's done! Hi @reedwm Can you please review this PR ? Thank you!
This PR has added ROCm for xla_call_module_test. 

/cc  @cheshire