diff --git a/Seq2SeqSharp/Tools/ComputeGraphTensor.cs b/Seq2SeqSharp/Tools/ComputeGraphTensor.cs
index 49d0110..fcbd7a1 100644
--- a/Seq2SeqSharp/Tools/ComputeGraphTensor.cs
+++ b/Seq2SeqSharp/Tools/ComputeGraphTensor.cs
@@ -102,7 +102,7 @@ public IWeightFactory GetWeightFactory()
 
         public IComputeGraph CreateSubGraph(string name)
         {
-            ComputeGraphTensor subGraph = new ComputeGraphTensor(m_weightTensorFactory, m_deviceId, m_needsBackprop, m_backprop, isSubGraph: true);
+            ComputeGraphTensor subGraph = new ComputeGraphTensor(m_weightTensorFactory, m_deviceId, m_needsBackprop, m_backprop, isSubGraph: true, saveGPUMemoryMode:m_saveGPUMemoryMode);
             return subGraph;
         }