network = nn.Sequential( nn.Linear(2, 5), nn.Tanh(), nn.Linear(5,1) ) opt = op.SGD(network.parameters(), lr=0.001) mse = nn.MSELoss() def train(model, opt, loss, X,y, epochs=1): for epoch in range(epochs): y2 = model(X) loss1 = loss(y2, y) loss1.backward() opt.step() opt.zero_grad() pass pass train(network, opt, mse, X, y, 10000)