def forward(self, x):
    out = self.embedding(x).permute(0, 2, 3, 1).view(x.size(0), (- 1), self.dims)
    for layer in self.blocks:
        out = layer(out)
    out = self.affine(out).mean(dim=1)
    return self.out(out)
def forward(self, x):
    out = self.embedding(x).permute(0, 2, 3, 1).view(x.size(0), (- 1), self.dims)
    for layer in self.blocks:
        out = layer(out)
    out = self.affine(out).mean(dim=1)
    return self.out(out)