Oh, I think I had nn.Linear there before, so it was 240 size in and 240 size out. I messed the example up because of that. It was supposed to be just 240, my bad!
I tried what you have suggested, and it looks like the following now:
import torch 
import torch.nn as nn 
import numpy as np

class SomeClass(nn.Module):

        def __init__(self):
             super().__init__()

             # create a learnable weight vector
             self.my_weights= nn.Parameter(torch.rand(1, 1, 1, 240))  # I need the vector to be learned weights
             nn.init.xavier_normal_(self.my_weights)
   
        def forward(q):

             # q = np.random.randn(150, 96, 96, 240)
             # size: [batch_size, sentence_length, sentence_length, weight_size]
             
            k = (q*self.my_weights)
            k = k.sum(3)

But I am very quickly getting an out of memory error on the multiplication part:
    k = (q*self.my_weights)
RuntimeError: CUDA error: out of memory

The memory jumps to over 6Gb VRAM (my limit on Nvidia 980Ti), while it usually only needs about 2-3 Gb. Any ideas about this? Thanks.
Also, what does the sum(3) actually do in this case? I hope the last column doesn’t just disappear, I need it applied to the previous two (as shown in the drawing I provided).