Upload folder using huggingface_hub

long_term_attention_gibbs.py

@@ -309,7 +309,7 @@ class LongTermAttention(nn.Module):
             self.queries = query.view(batch_size,qlen,self.n_head,self.d_head).transpose(1,2) # [B,h,q,d]
             self.keys = keys.view(batch_size,self.attn_num_basis,self.n_head,self.d_head).transpose(1,2) # [B,h,N,d]
             self.values = values.view(batch_size,self.attn_num_basis,self.n_head,self.d_head).transpose(1,2) # [B, h, q, N]
-            context = self.expected_value(self.score)  # Shape [1, 32, 768]
+            context = self.expected_value(self.score)  # Shape [B, n_head, qlen, d_head]
 
-            return context.contiguous().transpose(1,2).reshape(1, qlen, -1)
+            return context.contiguous().transpose(1,2).reshape(batch_size, qlen, -1)