diff --git a/.gitignore b/.gitignore
index c4ba141..fb0887d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,2 +1,3 @@
 .idea
+*.pyc
 datasets/mnist
diff --git a/1_linear_regression.py b/1_linear_regression.py
index b10aacd..4c754ae 100644
--- a/1_linear_regression.py
+++ b/1_linear_regression.py
@@ -17,7 +17,7 @@ def train(model, loss, optimizer, x, y):
     optimizer.zero_grad()
 
     # Forward
-    fx = model.forward(x.view(len(x), 1))
+    fx = model.forward(x.view(len(x), 1)).squeeze()
     output = loss.forward(fx, y)
 
     # Backward
@@ -26,7 +26,7 @@ def train(model, loss, optimizer, x, y):
     # Update parameters
     optimizer.step()
 
-    return output.data[0]
+    return output.item()
 
 
 def main():
@@ -35,7 +35,7 @@ def main():
     Y = 2 * X + torch.randn(X.size()) * 0.33
 
     model = build_model()
-    loss = torch.nn.MSELoss(size_average=True)
+    loss = torch.nn.MSELoss(reduction='elementwise_mean')
     optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
     batch_size = 10