From 161ecddafe2780d9aec600f96cf86cc1d0305170 Mon Sep 17 00:00:00 2001
From: "Godavarthy Surya, Anusha" <Anusha.GodavarthySurya@amd.com>
Date: Wed, 14 May 2025 20:52:58 +0530
Subject: [PATCH] SWDEV-519722 - Update test to work on MGpu machines without
 P2P access (#145)

---
 catch/unit/atomics/arithmetic_common.hh | 1 +
 1 file changed, 1 insertion(+)

diff --git a/catch/unit/atomics/arithmetic_common.hh b/catch/unit/atomics/arithmetic_common.hh
index ff6096d18f..4fb2111d74 100644
--- a/catch/unit/atomics/arithmetic_common.hh
+++ b/catch/unit/atomics/arithmetic_common.hh
@@ -446,6 +446,7 @@ void TestCore(const TestParams& p) {
 
   // Copy results back to Host
   for (auto i = 0u; i < p.num_devices; ++i) {
+    HIP_CHECK(hipSetDevice(i));
     const auto device_offset = i * p.kernel_count * p.ThreadCount();
     HIP_CHECK(hipMemcpy(old_vals.data() + device_offset, old_vals_devs[i].ptr(),
                         old_vals_alloc_size, hipMemcpyDeviceToHost));