Method forward

ML/src/python/neuralforge/nn/attention.py:125–128

(self, x, mask=None)

        self.norm = nn.LayerNorm(embed_dim)
    
    def forward(self, x, mask=None):
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

class VisionTransformerBlock(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_channels=3, embed_dim=768,