Indexing for plotting the weights

I want to plot the weights over time. lets say I have a model likenn.Linear(2, 2) I would like to plot all the weights in time. I plan to transform the weights into list of size four, so that I can just call as the pseudo code below

for i in weights.size()[0]:
   for j in weights.size()[1]:
       grid.plot(weights[i,j])

where weights[i,j] -> list - the values of weights[i,j] collected during training.

Or is there any other feature already like this exist in pytorch?

weights.data.view(-1).tolist() maybe helpful

1 Like