fix Rocm and Python lint

microsoft · Nov 29, 2023 · 8fddb03 · 8fddb03
1 parent f81c261
commit 8fddb03
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 2 deletions.
diff --git a/onnxruntime/core/providers/rocm/rocm_execution_provider.cc b/onnxruntime/core/providers/rocm/rocm_execution_provider.cc
@@ -320,8 +320,6 @@ Status ROCMExecutionProvider::Sync() const {
 }
 
 Status ROCMExecutionProvider::OnRunStart() {
-  // always set ROCM device when session::Run() in case it runs in a worker thread
-  HIP_RETURN_IF_ERROR(hipSetDevice(GetDeviceId()));
   return Status::OK();
 }
 

diff --git a/onnxruntime/core/providers/rocm/rocm_stream_handle.cc b/onnxruntime/core/providers/rocm/rocm_stream_handle.cc
@@ -181,6 +181,7 @@ void RegisterRocmStreamHandles(IStreamCommandHandleRegistry& stream_handle_regis
   stream_handle_registry.RegisterWaitFn(device_type, OrtDevice::CPU, WaitRocmNotificationOnHost);
   if (!use_existing_stream)
     stream_handle_registry.RegisterCreateStreamFn(device_type, [cpu_allocator, release_cpu_buffer_on_rocm_stream](const OrtDevice& device) {
+      HIP_CALL_THROW(hipSetDevice(device.Id()));
       hipStream_t stream = nullptr;
       HIP_CALL_THROW(hipStreamCreateWithFlags(&stream, hipStreamNonBlocking));
       return std::make_unique<RocmStream>(stream, device, cpu_allocator, release_cpu_buffer_on_rocm_stream, true, nullptr, nullptr);

diff --git a/onnxruntime/test/python/onnxruntime_test_python.py b/onnxruntime/test/python/onnxruntime_test_python.py
@@ -1740,5 +1740,6 @@ def test_iobinding_multiple_devices(self):
                 sessions[i].run_with_iobinding(binding)
                 binding.synchronize_outputs()
 
+
 if __name__ == "__main__":
     unittest.main(verbosity=1)