Sorry, I don’t think I fully understand what you are suggesting.
So, instead of
self.my_weights= nn.Parameter(torch.FloatTensor(240, 240)
I should somehow make it into:
self.my_weights= nn.Parameter(torch.FloatTensor([1,1,1,240], [1,1,1,240])) <-- that doesn’t seem to work in PyTorch. How can I initialize parameters with multiple dimensions properly?
So, based on my code snippet in the original post, I would do something like the following:
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.FloatTensor([1,1,1,240], [1,1,1,240]))
             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).sum(3)


Is this what you are suggesting or am I misunderstanding you? Also, thanks for looking into this!