Skip to content

Commit

Permalink
Bug fix for size mismatch issue occurred when using recent PyTorch.
Browse files Browse the repository at this point in the history
  • Loading branch information
vinhkhuc committed Oct 3, 2018
1 parent cbf48ed commit 73260e3
Show file tree
Hide file tree
Showing 2 changed files with 4 additions and 3 deletions.
1 change: 1 addition & 0 deletions .gitignore
Original file line number Diff line number Diff line change
@@ -1,2 +1,3 @@
.idea
*.pyc
datasets/mnist
6 changes: 3 additions & 3 deletions 1_linear_regression.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -26,7 +26,7 @@ def train(model, loss, optimizer, x, y):
# Update parameters
optimizer.step()

return output.data[0]
return output.item()


def main():
Expand All @@ -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

Expand Down

0 comments on commit 73260e3

Please sign in to comment.