From 6a0bd4a47cb2f497354ce63ec5c2f230fbf100cf Mon Sep 17 00:00:00 2001
From: fruitea <aetiurf@gmail.com>
Date: Sun, 11 Aug 2024 14:16:52 -0700
Subject: [PATCH] chore: minor

---
 include/flexflow/utils/memory_allocator.h | 2 +-
 src/utils/memory_allocator.cc             | 7 ++++---
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/include/flexflow/utils/memory_allocator.h b/include/flexflow/utils/memory_allocator.h
index 2b84f4f2ae..7091b159b2 100644
--- a/include/flexflow/utils/memory_allocator.h
+++ b/include/flexflow/utils/memory_allocator.h
@@ -62,7 +62,7 @@ class MemoryAllocator {
   size_t instance_total_size, instance_allocated_size;
 };
 
-Legion::Memory get_proc_mem(Legion::Processor proc);
+Legion::Memory get_proc_mem(Legion::Machine machine, Legion::Processor proc);
 
 }; // namespace FlexFlow
 
diff --git a/src/utils/memory_allocator.cc b/src/utils/memory_allocator.cc
index 4e95ba0eb5..792b529b95 100644
--- a/src/utils/memory_allocator.cc
+++ b/src/utils/memory_allocator.cc
@@ -19,10 +19,11 @@ namespace FlexFlow {
 
 using Legion::Machine;
 using Legion::Memory;
+using Legion::Processor;
 
-Memory get_proc_mem(Legion::Processor proc) {
+Memory get_proc_mem(Machine machine, Processor proc) {
     // First try to allocate a managed memory (cudaMallocManaged)
-    Memory proc_mem = Machine::MemoryQuery(Machine::get_machine())
+    Memory proc_mem = Machine::MemoryQuery(machine)
                         .only_kind(Memory::GPU_MANAGED_MEM)
                         .best_affinity_to(proc)
                         .first();
@@ -30,7 +31,7 @@ Memory get_proc_mem(Legion::Processor proc) {
       return proc_mem;
     }
     // If managed memory is not available, try to allocate a framebuffer memory
-    proc_mem = Machine::MemoryQuery(Machine::get_machine())
+    proc_mem = Machine::MemoryQuery(machine)
                         .only_kind(Memory::GPU_FB_MEM)
                         .best_affinity_to(proc)
                         .first();