1 files changed, 23 insertions, 2 deletions
diff --git a/src/backends/cl/workloads/ClBaseWorkload.hpp b/src/backends/cl/workloads/ClBaseWorkload.hpp
index e74fc84f4f..03417e33ae 100644
--- a/src/backends/cl/workloads/ClBaseWorkload.hpp
+++ b/src/backends/cl/workloads/ClBaseWorkload.hpp
@@ -20,17 +20,38 @@ public:
     // Replace input tensor handle with the given TensorHandle and call Reconfigure()
     void ReplaceInputTensorHandle(ITensorHandle* tensorHandle, unsigned int slot) override
     {
+        ITensorHandle* backupHandle = this->m_Data.m_Inputs[slot];
         this->m_Data.m_Inputs[slot] = tensorHandle;
-        Reconfigure();
+        try
+        {
+            Reconfigure();
+        }
+        catch(armnn::UnimplementedException& e)
+        {
+            // Cannot reconfigure, revert the slot back and throw the exception.
+            this->m_Data.m_Inputs[slot] = backupHandle;
+            throw e;
+        }
     }
 
     // Replace output tensor handle with the given TensorHandle and call Reconfigure()
     void ReplaceOutputTensorHandle(ITensorHandle* tensorHandle, unsigned int slot) override
     {
+        ITensorHandle* backupHandle = this->m_Data.m_Outputs[slot];
         this->m_Data.m_Outputs[slot] = tensorHandle;
-        Reconfigure();
+        try
+        {
+            Reconfigure();
+        }
+        catch(armnn::UnimplementedException& e)
+        {
+            // Cannot reconfigure, revert the slot back and throw the exception.
+            this->m_Data.m_Inputs[slot] = backupHandle;
+            throw e;
+        }
     }
 
+protected:
     // Reconfigure the workload configuration. Throw armnn::UnimplementedException by default.
     virtual void Reconfigure()
     {