small bug fix in tnn

1 files changed, 6 insertions, 2 deletions

diff --git a/nerv/tnn/tnn.lua b/nerv/tnn/tnn.lua
index e7d4981..d527fe6 100644
--- a/nerv/tnn/tnn.lua
+++ b/nerv/tnn/tnn.lua
@@ -347,7 +347,9 @@ function TNN:net_propagate() --propagate according to feeds_now
     local feeds_now = self.feeds_now
     for t = 1, self.chunk_size do --some layer maybe do not have inputs from time 1..chunk_size
         for id, ref in pairs(self.layers) do
-            self:propagate_dfs(ref, t)
+            if #ref.dim_in > 0 then --some layer is just there(only to save some parameter)
+                self:propagate_dfs(ref, t)
+            end
         end
     end
     for t = 1, self.chunk_size do
@@ -472,7 +474,9 @@ function TNN:net_backpropagate(do_update) --propagate according to feeds_now
     local feeds_now = self.feeds_now
     for t = 1, self.chunk_size do --some layer maybe do not have outputs from time 1..chunk_size
         for id, ref in pairs(self.layers) do
-            self:backpropagate_dfs(ref, t, do_update)
+            if #ref.dim_out > 0 then --some layer is just there(only to save some parameter)
+                self:backpropagate_dfs(ref, t, do_update)
+            end
         end
     end
     for t = 1, self.chunk_size do