From 2fc776c0f9e62a9f795ed1a0552c84fb1ac644db Mon Sep 17 00:00:00 2001
From: Jong Wook Kim <jongwook@nyu.edu>
Date: Thu, 21 Jul 2022 12:51:00 -0700
Subject: [PATCH] simplified flatten

---
 clip/model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/clip/model.py b/clip/model.py
index f264b63..808bf16 100644
--- a/clip/model.py
+++ b/clip/model.py
@@ -66,7 +66,7 @@ class AttentionPool2d(nn.Module):
         self.num_heads = num_heads
 
     def forward(self, x):
-        x = x.reshape(x.shape[0], x.shape[1], x.shape[2] * x.shape[3]).permute(2, 0, 1)  # NCHW -> (HW)NC
+        x = x.flatten(start_dim=2).permute(2, 0, 1)  # NCHW -> (HW)NC
         x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0)  # (HW+1)NC
         x = x + self.positional_embedding[:, None, :].to(x.dtype)  # (HW+1)NC
         x, _ = F.multi_head_attention_forward(