Update modeling_nort5.py

modeling_nort5.py

@@ -62,7 +62,7 @@ class Decoder(nn.Module):
         self_relative_embedding = self.self_relative_embedding()
         cross_relative_embedding = self.cross_relative_embedding()
 
-        if past_key_values is not None:
+        if past_key_values is None:
             autoreg_mask = torch.triu(
                 torch.full((x.size(0), x.size(0)), True, device=x.device),
                 diagonal=1
@@ -259,12 +259,12 @@ class Attention(nn.Module):
 
         if past_key_value is not None:
             if not self.is_cross_attention:
-                key = torch.cat([past_key_value[0], key], dim=1)
-                value = torch.cat([past_key_value[1], value], dim=1)
+                key = torch.cat([past_key_value[0].flatten(0, 1), key], dim=1)
+                value = torch.cat([past_key_value[1].flatten(0, 1), value], dim=1)
                 key_len = key.size(1)
             elif past_key_value[0].size(1) == kv.size(0):
-                key = past_key_value[0]
-                value = past_key_value[1]
+                key = past_key_value[0].flatten(0, 1)
+                value = past_key_value[1].flatten(0, 1)
 
         if self.position_indices.size(0) < max(query_len, key_len):
             position_indices = torch.arange(max(query_len, key_len), dtype=torch.long).unsqueeze(1) \
@@ -306,7 +306,10 @@ class Attention(nn.Module):
         context = self.post_layer_norm(context)
         context = self.dropout(context)
 
-        return context, attention_probs.detach(), (key.detach(), value.detach())
+        key = key.detach().unflatten(0, (-1, self.num_heads))
+        value = value.detach().unflatten(0, (-1, self.num_heads))
+
+        return context, attention_probs.detach(), (key, value)
 
 
 class WordEmbedding(nn.Module):
@@ -662,9 +665,7 @@ class NorT5ForConditionalGeneration(NorT5Model):
             reordered_layer_past_states = ()
             for layer_past_state in layer_past_states:
                 # need to set correct `past` for each of the four key / value states
-                layer_past_state = layer_past_state.unflatten(0, (-1, self.config.num_attention_heads))
                 layer_past_state = layer_past_state.index_select(0, beam_idx.to(layer_past_state.device))
-                layer_past_state = layer_past_state.flatten(0, 1)
                 reordered_layer_past_states = reordered_layer_past_states + (layer_past_state,)
 
             assert reordered_layer_past_states[0].shape == layer_past_states[0].shape