def __init__(...):
    self.head = nn.Linear(64, 2 * 512)


def forward(....):
    A = self.head(E).reshape(-1, 512, 2).softmax(dim=-1)[:, :, 1]

But if you are doing a 2-element softmax, you might as well just use a 1-element output + sigmoid instead. I.e.,
def __init__(...):
    self.head = nn.Linear(64, 512)


def forward(....):
    A = self.head(E).sigmoid()