So, changing the weight to nn.Parameter should work, right?
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import numpy as np

class Setting(nn.Module):
    
    def __init__(self, A):
        
        super(Setting, self).__init__()
        
        self.A = A
        self.X = np.array(np.sum(A, axis=1))
        self.feature_len = 10

        self.L = 4
        self.n = 10
        self.z_dim = 8
        
        self.h = torch.empty(self.L-1,self.n, self.z_dim)
        
        self.h0 = torch.from_numpy(self.X)
        
        W0 = nn.parameter.Parameter(torch.FloatTensor((self.feature_len, self.z_dim)))
        
        self.h[0] = torch.empty(self.n, self.z_dim)

        for v in range(self.n):
            self.h[0][v] = F.relu(self.h0[v]*W0)
        
    def forward(self, A):
        
        d_u = np.empty([self.n, self.n])
        
        for v in range(self.n):
            d_u[v] = self.X[v]*self.A[v]

        d_u = torch.from_numpy(d_u)        
                
        h_n = torch.empty(self.L-1, self.n, self.z_dim)

        rnn = nn.GRUCell(self.z_dim, self.z_dim)

        H = F.normalize(self.h[0], p=2, dim=1)
        
        for l in range(1,self.L-1):
            
            h_n[l] = torch.mm(d_u, self.h[l-1].double())
            self.h[l] = rnn(self.h[l-1], h_n[l]).detach().float()
            self.h[l] = F.normalize(self.h[l], p=2, dim=1)
            H = torch.max(H, self.h[l])
   
        return H


So, I’ve updated the weight as Parameter, but I’m unclear on how to add W0 in the forward method. Also, note that the limits of the for loop are changed so that there is no overwriting.
So, my question is as follows:
I first calculate the h[0] using W and h0 (this is different from h[0]). Now further, I’ll use these values to calculate h[1], h[2], …, h[L-1]. Max value from each of this is updated is updated into H with

H = torch.max(H, self.h[l])

This returned value H is compared with H1 (another already known embedding).
So, to make H and H1 embeddings equal, weight W0 is basic the trainable parameter. How to do that and what are the modifications required in the code?
Thanks a lot for your time and help.