Hi, I have the following tensor size q=[150, 96, 96, 240] (here: [batch_size, sentence_length, sentence_length, weight_size]), I need to multiply it with a weight vector of size [240] and either reshape or use convolution to make q of size [150, 96, 96], basically combining the last [96, 96, 240] into size [96, 96]. I have a short code snippet, but I don’t know how to properly reshape, and if doing convolution, in this case, makes more sense (and if so, how to do it).
Any help appreciated.

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(240)
             nn.init.xavier_normal_(self.my_weights)
   
        def forward(q):

             # q = np.random.randn(150, 96, 96, 240)
             # a = torch.from_numpy(q)  
             # size: [batch_size, sentence_length, sentence_length, weight_size]
             
            # multiply the weight vector with the last column of the tensor, is this correct?
            k= torch.matmul(q, self.my_weights)
            
            # reshape or do a 2d Convolution
            v = TODO(k)  [150, 96, 96, 240] -->  [150, 96, 96, 1] --> [150, 96, 96]


Visual representation I drew up: https://imgur.com/a/yuIc3t4 4