From 885d95d18b973710efc87cc0b62754266df2e5e1 Mon Sep 17 00:00:00 2001
From: Anton Oresten <anton.oresten42@gmail.com>
Date: Fri, 13 Dec 2024 17:57:26 +0100
Subject: [PATCH] Fix device

---
 src/distributed/public_api.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/distributed/public_api.jl b/src/distributed/public_api.jl
index d5d10e42a4..51a1dd9b9a 100644
--- a/src/distributed/public_api.jl
+++ b/src/distributed/public_api.jl
@@ -132,7 +132,7 @@ Backend Agnostic API to perform an allreduce operation on the given buffer `send
 workers.
 """
 function allreduce!(backend::AbstractFluxDistributedBackend, sendrecvbuf, op::F) where {F}
-    return __allreduce!(backend, sendrecvbuf, op, gpu_device())
+    return __allreduce!(backend, sendrecvbuf, op, get_device(sendrecvbuf))
 end
 
 function allreduce!(