Optimize memory usage for EltMulMulAdd operatior

zhongkaifu · Sep 20, 2023 · 71e9372 · 71e9372
1 parent caff84b
commit 71e9372
Showing 1 changed file with 20 additions and 4 deletions.
diff --git a/Seq2SeqSharp/Tools/ComputeGraphTensor.cs b/Seq2SeqSharp/Tools/ComputeGraphTensor.cs
@@ -495,10 +495,26 @@ void backward()
                 m_backprop.Add(backward);
 
                 // These tensors' weights will be used during back-propogation, so we unbind them from the computing graph
-                m1.UnbindFromComputeGraph();
-                m2.UnbindFromComputeGraph();
-                m3.UnbindFromComputeGraph();
-                m4.UnbindFromComputeGraph();
+
+                if (m2.NeedGradient)
+                {
+                    m1.UnbindFromComputeGraph();
+                }
+
+                if (m1.NeedGradient)
+                {
+                    m2.UnbindFromComputeGraph();
+                }
+
+                if (m4.NeedGradient)
+                {
+                    m3.UnbindFromComputeGraph();
+                }
+
+                if (m3.NeedGradient)
+                {
+                    m4.UnbindFromComputeGraph();
+                }
             }