What’s the proper way to do univariate optimization with PyTorch?

Currently, what I do is something like the following which correctly finds the mean of given sample pts as 2.5.

```
import torch
from torch.autograd import Variable
mean = Variable(torch.ones(1, 1), requires_grad=True)
criterion = torch.nn.MSELoss()
optimizer = torch.optim.SGD([mean], lr=0.01, momentum=0.9)
pts = torch.Tensor([1, 2, 3, 4])
for _ in range(10**3):
optimizer.zero_grad()
loss = criterion(mean, pts)
loss.backward()
optimizer.step()
print(mean.data)
```