Sure thing. The more or less full code looks like this (with random data):
import torch 
import torch.nn as nn 
import numpy as np

class SomeClass(nn.Module):

        def __init__(self):
             super().__init__()
            
             self.relu = nn.RReLU()
             self.conv = nn.Conv2d(240, kernel_size=1, out_channels=1)
   
        def forward(q, k, values):
             
             # q = k = values
             q = np.random.randn(150, 96, 120)
             k = np.random.randn(150, 96, 120)
             
             # now transform both into one tensor of size: [150, 96, 96, 240]
             # not specifically related to the question above though
             a = torch.from_numpy(q)
             b = torch.from_numpy(k)
             bs, s, v = a.size()
             a_ = a.repeat(1, 1, s).view(bs, s*s, v)
             b_ = b.repeat(1, s, 1)
             concat_vec = torch.stack((a_, b_), 2).view(bs, s, s, -1)
             # concat_vec size: 
             # [batch_size, sentence_length, sentence_length, weight_size] = [150, 96, 96, 240]
             
             # now transform into [150, 96, 96]
             attention = self.conv(concat_vec.permute((0, 3, 1, 2))).squeeze()

             attention = self.relu(attention)
             output = torch.bmm(attention, values)

             return output

In case this helps, the code is the modification of the scaled-dot product attention from the Self-attention paper (https://arxiv.org/abs/1703.03130), the modification is the implementation of the concatenation weighting function from this paper: https://arxiv.org/abs/1711.07971 (Section 3.2).
Originally, the code looks like this without the modification and works just fine:
import torch 
import torch.nn as nn 
import numpy as np

class SomeClass(nn.Module):

        def __init__(self):
             super().__init__()
   
        def forward(q, k, values):
             
             attn = torch.bmm(q, k.transpose(1, 2))
             attn = self.softmax(attn)
             output = torch.bmm(attn, values)
             return output

Error I am getting (happens later on in the code that uses the transformed output):
File "...\lib\site-packages\torch\tensor.py", line 381, in __iter__
    raise TypeError('iteration over a 0-d tensor')
TypeError: iteration over a 0-d tensor