diff --git a/src/runtime/CL/CLTensor.cpp b/src/runtime/CL/CLTensor.cpp
index 9bbf926..a6d0cf7 100644
--- a/src/runtime/CL/CLTensor.cpp
+++ b/src/runtime/CL/CLTensor.cpp
@@ -23,15 +23,21 @@
  */
 #include "arm_compute/runtime/CL/CLTensor.h"
 
+#include "arm_compute/runtime/CL/CLRuntimeContext.h"
 #include "arm_compute/runtime/CL/CLScheduler.h"
 
 namespace arm_compute
 {
-CLTensor::CLTensor()
-    : _allocator(this)
+CLTensor::CLTensor(IRuntimeContext *ctx)
+    : _allocator(this, static_cast<CLRuntimeContext *>(ctx)), _ctx(static_cast<CLRuntimeContext *>(ctx))
 {
 }
 
+CLRuntimeContext *CLTensor::context()
+{
+    return _ctx;
+}
+
 TensorInfo *CLTensor::info() const
 {
     return &_allocator.info();
@@ -59,12 +65,12 @@
 
 void CLTensor::map(bool blocking)
 {
-    ICLTensor::map(CLScheduler::get().queue(), blocking);
+    ICLTensor::map(_ctx == nullptr ? CLScheduler::get().queue() : _ctx->gpu_scheduler()->queue(), blocking);
 }
 
 void CLTensor::unmap()
 {
-    ICLTensor::unmap(CLScheduler::get().queue());
+    ICLTensor::unmap(_ctx == nullptr ? CLScheduler::get().queue() : _ctx->gpu_scheduler()->queue());
 }
 
 uint8_t *CLTensor::do_map(cl::CommandQueue &q, bool blocking)
@@ -81,4 +87,4 @@
 {
     _allocator.set_associated_memory_group(memory_group);
 }
-} // namespace arm_compute
\ No newline at end of file
+} // namespace arm_compute