Hi @WeiHao97,
If I understand correctly you want to train your model w.r.t. B, the learning parameter matrix.
class Network(nn.Module):
   def __init__(self, shape):
       super().__init__()
       self.B = nn.Parameter(torch.zeros(shape))
       torch.nn.init.xavier_uniform_(self.B) #or any other init method
   def forward(self, x, A):
       M = A*self.B
       return x @ M.t()

shape has to be (output_dim, input_dim).  input_dim being the dim of x.
Hope this is what you wanted to do.