(self, x, skip)
| 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()) |
| 631 | x = self.proj(x) |
| 632 | # print("x after proj in Tokensplit is ", x.mean(), x.std()) |
| 633 | x = rearrange(x, "... h w (nh nw e) -> ... (h nh) (w nw) e", nh=self.h, nw=self.w) |
| 634 | # print("x after rearange in Tokensplit is ", x.mean(), x.std()) |
| 635 | x = torch.lerp(skip, x, self.fac.to(x.dtype)) |
| 636 | # x = mp_sum_t(x, skip, self.fac.to(x.dtype)) |
| 637 | # print("x after lerp in Tokensplit is ", x.mean(), x.std()) |
| 638 | return x |
| 639 | |
| 640 | |
| 641 | class LocalCondProj(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected