remove erroneous formatting

This commit is contained in:
Penn 2022-07-21 11:53:46 -07:00
parent 8c0b98adfc
commit 5eb1828ead
1 changed files with 2 additions and 2 deletions

View File

@ -67,8 +67,8 @@ class AttentionPool2d(nn.Module):
def forward(self, x): 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.reshape(x.shape[0], x.shape[1], x.shape[2] * x.shape[3]).permute(2, 0, 1) # NCHW -> (HW)NC
x = torch.cat([x.mean(dim=0, keepdim=True), x], dim=0) # (HW+1)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 = x + self.positional_embedding[:, None, :].to(x.dtype) # (HW+1)NC
x, _ = F.multi_head_attention_forward( x, _ = F.multi_head_attention_forward(
query=x[:1], key=x, value=x, query=x[:1], key=x, value=x,
embed_dim_to_check=x.shape[-1], embed_dim_to_check=x.shape[-1],