From bbacc266cd43cc16a2e17dead7b9bc6011952460 Mon Sep 17 00:00:00 2001
From: tankya2 <tankya2@ihpc.a-star.edu.sg>
Date: Thu, 4 Jul 2024 16:10:43 +0800
Subject: [PATCH] correct missing mempool initialization

---
 src/qibotn/eval.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/src/qibotn/eval.py b/src/qibotn/eval.py
index 23624aa..5a012bb 100644
--- a/src/qibotn/eval.py
+++ b/src/qibotn/eval.py
@@ -320,6 +320,7 @@ def expectation_pauli_tn_MPI(qibo_circ, datatype, pauli_string_pattern, n_sample
     """
     from cuquantum import Network
     from mpi4py import MPI  # this line initializes MPI
+    import cuquantum.cutensornet as cutn
 
     root = 0
     comm = MPI.COMM_WORLD
@@ -329,6 +330,7 @@ def expectation_pauli_tn_MPI(qibo_circ, datatype, pauli_string_pattern, n_sample
     # Assign the device for each process.
     device_id = rank % getDeviceCount()
     cp.cuda.Device(device_id).use()
+    mempool = cp.get_default_memory_pool()
 
     # Perform circuit conversion
     if rank==0: