diff --git a/src/PIM/Compiler/PimCodeGen.cpp b/src/PIM/Compiler/PimCodeGen.cpp
index 3c29ad9..dba3664 100644
--- a/src/PIM/Compiler/PimCodeGen.cpp
+++ b/src/PIM/Compiler/PimCodeGen.cpp
@@ -140,7 +140,7 @@ void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) {
   });
 
   funcOp.walk([&](memref::AllocOp allocOp) {
-    if (!allocOp->getParentOfType<pim::PimCoreOp>())
+    if (!allocOp->getParentOfType<pim::PimCoreOp>() && !allocOp->getParentOfType<pim::PimCoreBatchOp>())
       gatherMemEntry(allocOp.getResult());
   });