Instead of modifying inplace, can you try creating a clone of the object. Something like
yhat = net(inputs) 
yhat_clone = yhat.clone()
yhat_clone[0] = yhat[1]   # here
yhat_clone[-1] = yhat[-2] # here
loss = free_energy(yhat_clone, x, matrix, rho, lambd) 
loss_points.append(loss.item())