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();