change print

HNR1 · Jun 12, 2023 · 2acf9b7 · 2acf9b7
1 parent d6a9905
commit 2acf9b7
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tomesd/patch.py b/tomesd/patch.py
@@ -127,9 +127,11 @@ def forward(
                 norm_hidden_states = (
                     self.norm2(hidden_states, timestep) if self.use_ada_layer_norm else self.norm2(hidden_states)
                 )
+                print('n_hidden', norm_hidden_states[0,0,:6])
+
                 # (4) ToMe m_c
                 norm_hidden_states = m_c(norm_hidden_states)
-                print('n_hidden', norm_hidden_states[0,0,:6])
+
                 # 2. Cross-Attention
                 attn_output = self.attn2(
                     norm_hidden_states,