Method __init__

ML/src/python/neuralforge/nn/modules.py:168–170

(self, kernel_size=7)

class SpatialAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super().__init__()
        self.conv = nn.Conv2d(2, 1, kernel_size, padding=kernel_size // 2)
    
    def forward(self, x):
        avg_out = torch.mean(x, dim=1, keepdim=True)