4 files changed, 13 insertions, 8 deletions
diff --git a/src/runtime/CL/CLTensor.cpp b/src/runtime/CL/CLTensor.cpp
index 732689e7ec..9bbf926b58 100644
--- a/src/runtime/CL/CLTensor.cpp
+++ b/src/runtime/CL/CLTensor.cpp
@@ -25,8 +25,8 @@
 
 #include "arm_compute/runtime/CL/CLScheduler.h"
 
-using namespace arm_compute;
-
+namespace arm_compute
+{
 CLTensor::CLTensor()
     : _allocator(this)
 {
@@ -76,3 +76,9 @@ void CLTensor::do_unmap(cl::CommandQueue &q)
 {
     _allocator.unmap(q, buffer());
 }
+
+void CLTensor::associate_memory_group(arm_compute::IMemoryGroup *memory_group)
+{
+    _allocator.set_associated_memory_group(memory_group);
+}
+} // namespace arm_compute
+\ No newline at end of file
diff --git a/src/runtime/CL/CLTensorAllocator.cpp b/src/runtime/CL/CLTensorAllocator.cpp
index 72b5854c5c..60e1ca8df6 100644
--- a/src/runtime/CL/CLTensorAllocator.cpp
+++ b/src/runtime/CL/CLTensorAllocator.cpp
@@ -25,7 +25,6 @@
 
 #include "arm_compute/core/Error.h"
 #include "arm_compute/core/TensorInfo.h"
-#include "arm_compute/runtime/CL/CLMemoryGroup.h"
 #include "arm_compute/runtime/CL/CLScheduler.h"
 
 namespace arm_compute
@@ -104,8 +103,8 @@ void populate_quantization_info(CLFloatArray &scale, CLInt32Array &offset, const
 }
 } // namespace
 
-CLTensorAllocator::CLTensorAllocator(CLTensor *owner)
-    : _associated_memory_group(nullptr), _memory(), _mapping(nullptr), _owner(owner), _scale(), _offset()
+CLTensorAllocator::CLTensorAllocator(IMemoryManageable *owner)
+    : _owner(owner), _associated_memory_group(nullptr), _memory(), _mapping(nullptr), _scale(), _offset()
 {
 }
 
@@ -143,7 +142,7 @@ void CLTensorAllocator::allocate()
     }
     else
     {
-        _associated_memory_group->finalize_memory(_owner, _memory, info().total_size());
+        _associated_memory_group->finalize_memory(_owner, _memory, info().total_size(), alignment());
     }
 
     // Allocate and fill the quantization parameter arrays
@@ -178,7 +177,7 @@ Status CLTensorAllocator::import_memory(cl::Buffer buffer)
     return Status{};
 }
 
-void CLTensorAllocator::set_associated_memory_group(CLMemoryGroup *associated_memory_group)
+void CLTensorAllocator::set_associated_memory_group(IMemoryGroup *associated_memory_group)
 {
     ARM_COMPUTE_ERROR_ON(associated_memory_group == nullptr);
     ARM_COMPUTE_ERROR_ON(_associated_memory_group != nullptr && _associated_memory_group != associated_memory_group);
diff --git a/src/runtime/CL/functions/CLDeconvolutionLayerUpsample.cpp b/src/runtime/CL/functions/CLDeconvolutionLayerUpsample.cpp
index 63a45aae8e..eaf7c66083 100644
--- a/src/runtime/CL/functions/CLDeconvolutionLayerUpsample.cpp
+++ b/src/runtime/CL/functions/CLDeconvolutionLayerUpsample.cpp
@@ -26,6 +26,7 @@
 #include "arm_compute/core/CL/OpenCL.h"
 #include "arm_compute/core/Utils.h"
 #include "arm_compute/runtime/CL/CLScheduler.h"
+#include "arm_compute/runtime/CL/CLTensor.h"
 
 namespace arm_compute
 {
diff --git a/src/runtime/CL/functions/CLSoftmaxLayer.cpp b/src/runtime/CL/functions/CLSoftmaxLayer.cpp
index 7e41dba8ab..73add97ef1 100644
--- a/src/runtime/CL/functions/CLSoftmaxLayer.cpp
+++ b/src/runtime/CL/functions/CLSoftmaxLayer.cpp
@@ -30,7 +30,6 @@
 #include "arm_compute/core/Types.h"
 #include "arm_compute/core/Utils.h"
 #include "arm_compute/core/utils/misc/ShapeCalculator.h"
-#include "arm_compute/runtime/CL/CLMemoryGroup.h"
 #include "arm_compute/runtime/CL/CLScheduler.h"
 
 namespace arm_compute