Update:
This also doesn’t work
AB = torch.zeros(bsz, num_classes).cuda()
for j in range(0, embed_dim, jbatch):
    AB += torch.abs(A[:,j:j+jbatch].unsqueeze(-1) - B[j:j+jbatch,:]).sum(dim=-2)

where A is my model’s (bsz, embed_dim) output, B is my (embed_dim, num_classes) embedding layer, AB is my (bsz, num_classes) logits, and jbatch is a batching parameter to speed things up depending on available GPU memory.
as I loop through j, the allocated GPU memory keeps increasing by AB and I run out of memory quite fast