From 61b4e716ca6b58d82f8afe75472c5455e4525519 Mon Sep 17 00:00:00 2001 From: Fredo Guan Date: Wed, 18 Jan 2023 12:30:37 -0800 Subject: [PATCH] Update metaformers.py --- timm/models/metaformers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/timm/models/metaformers.py b/timm/models/metaformers.py index d0c2b401..fd992209 100644 --- a/timm/models/metaformers.py +++ b/timm/models/metaformers.py @@ -180,7 +180,7 @@ class RandomMixing(nn.Module): data=torch.softmax(torch.rand(num_tokens, num_tokens), dim=-1), requires_grad=False) ''' - self.random_matrix = torch.softmax(torch.rand(num_tokens, num_tokens)) + self.random_matrix = torch.softmax(torch.rand(num_tokens, num_tokens), dim=-1) def forward(self, x): B, H, W, C = x.shape x = x.reshape(B, H*W, C)