Hello, I’m trying to implement a variation on a linear layer:

More specifically:
I have a model with an output (bsz x embed_dim).
I have an embed_matrix (embed_dim x num_classes)
I am trying to define an operation that takes in output and embed_matrix, and returns logits (bsz x num_classes)
With a linear layer, each element (i,j) in the logits corresponds to cosine similarity between the i’th embedding output by my model, and the j’th embedding in my embed_matrix.
I’d like to create a custom linear layer where instead of cosine similarity, each element corresponds to the L1 distance between two embeddings.
I’ve tried:
torch.abs(B - A.unsqueeze(-1)).sum(dim=-2)

however, this causes the memory to explode when a Linear layer handles the same matrices just fine.
Is there a way for me to compute this output with a similar memory/time as a linear layer?
Thank you!