Update MT.py

MT.py

@@ -285,10 +285,11 @@ class FeatureTransformer(nn.Module):
         for i in range(self.num_layers):
             value, att, attn_viz = self.layers(att=att, value=value, shape=[h, w], iteration=i)
             value_decode = self.normalize(torch.square(self.re_proj(value)))  # map to motion energy, Do use normalization here
-        
-        attn_viz_list.append(attn_viz.reshape(b, h, w, h, w))
-        attn_list.append(att.view(b, h, w, c).permute(0, 3, 1, 2).contiguous())
-        feature_list.append(value_decode.view(b, h, w, c).permute(0, 3, 1, 2).contiguous())
+            if i % 2 == 0:
+                attn_viz_list.append(attn_viz.reshape(b, h, w, h, w))
+                attn_list.append(att.view(b, h, w, c).permute(0, 3, 1, 2).contiguous())
+                feature_list.append(value_decode.view(b, h, w, c).permute(0, 3, 1, 2).contiguous())
+            
         return feature_list, attn_list, attn_viz_list
 
     def forward_save_mem(self, feature0, add_position_embedding=True):