How to manually adjust parameters during training

How to manually adjust parameters during training? I tried to manually adjust the parameter values, but the gradient calculation was not correct in the backward of the next epoch. The parameters that I reset to zero will still be updated, but this cannot happen because the result obtained by the Hadema product of 0 elements should not have a gradient.

The optimizer is Adam.

I try to reproduce a minimal script that reproduces the problem:

self.user_embedding = torch.nn.Embedding(
     num_embeddings=len(self.user_map), embedding_dim=self.latent_dim
)
self.item_embedding = torch.nn.Embedding(
    num_embeddings=len(self.item_map), embedding_dim=self.latent_dim
)
self.adj_mat = self.get_adj_mat() # 0,1 adjacency matrix, does not contain gradients
self.adj_mask_train = torch.nn.Parameter(self.generate_adj_mask(self.generate_daj_mat()),requires_grad=True)
# 0,1 adjacency matrix mask, contain gradients

optimizer = optim.Adam(params, lr=learning_rate, weight_decay=weight_decay)
optimizer.zero_grad()

user_all_embeddings, item_all_embeddings = self.forward()
u_embeddings = user_all_embeddings[user]
pos_embeddings = item_all_embeddings[pos_item]
neg_embeddings = item_all_embeddings[neg_item]
pos_scores = torch.mul(u_embeddings, pos_embeddings).sum(dim=1)
neg_scores = torch.mul(u_embeddings, neg_embeddings).sum(dim=1)
mf_loss = self.mf_loss(pos_scores, neg_scores)

mf_loss.backward()
optimizer.step()



def forward:
    user_embeddings = self.user_embedding.weight
    item_embeddings = self.item_embedding.weight
    ego_embeddings = torch.cat([user_embeddings, item_embeddings], dim=0)
    embeddings_list = [ego_embeddings]
    if some condition: # prune the adj_mask_train
          with torch.no_grad():
            adj_mask_tensor = self.adj_mask_train.flatten()
            nonzero = torch.abs(adj_mask_tensor) > 0
            adj_mask = adj_mask_tensor[nonzero]  
            adj_total = adj_mask.shape[0]
            adj_y, adj_i = torch.sort(adj_mask.abs())
            adj_thre_index = int(adj_total * prune_percent)
            adj_thre = adj_y[adj_thre_index]

            abs_values = torch.abs(self.adj_mask_train)
            mask = abs_values >= adj_thre
            self.adj_mask_train.data[mask] = 1
            self.adj_mask_train.data[~mask] = 0
    dense_tensor = self.adj_mat
    adj = torch.mul(dense_tensor, self.adj_mask_train)
    self.norm_adj_matrix = self.torch_normalize_adj(adj)
    for layer_idx in range(self.n_layers):
        all_embeddings = torch.mm(self.norm_adj_matrix, all_embeddings)
        embeddings_list.append(all_embeddings)
    lightgcn_all_embeddings = torch.stack(embeddings_list, dim=1)
    lightgcn_all_embeddings = torch.mean(lightgcn_all_embeddings, dim=1)
    user_all_embeddings, item_all_embeddings = torch.split(
                lightgcn_all_embeddings, [self.n_users, self.n_items]
            )
    return user_all_embeddings, item_all_embeddings

        

Please don’t mind asking me for more details, I really need your help, thank you very much!
former question(the same):Manually modified parameter during trainning, gradients are not being correctly updated. Help