(self, in_features, out_features, patch_size=(2, 2))
| 620 | |
| 621 | class TokenSplit(nn.Module): |
| 622 | def __init__(self, in_features, out_features, patch_size=(2, 2)): |
| 623 | super().__init__() |
| 624 | self.h = patch_size[0] |
| 625 | self.w = patch_size[1] |
| 626 | self.proj = apply_wd(Linear(in_features, out_features * self.h * self.w, bias=False)) |
| 627 | self.fac = nn.Parameter(torch.ones(1) * 0.5) |
| 628 | |
| 629 | def forward(self, x, skip): |
| 630 | # print("x before proj in Tokensplit is ", x.mean(), x.std()) |