commit cb66e9cf78e9b4dbe8323a263ecf00f8eda27199
author NC Cullen <nicholas.c.cullen.th@dartmouth.edu> Wed Apr 12 23:59:12 2017 -0400
committer Soumith Chintala <soumith@gmail.com> Wed Apr 12 20:59:12 2017 -0700
tree 3a6adc86a594cd69a029d25d2bbdf47bba6839e8
parent c85288308658b3d5ca631366828e3a5befc94691

torch.diag bug fix (#1251)

test/test_autograd.py

diff --git a/test/test_autograd.py b/test/test_autograd.py
index a20f25d..3a1677b 100644
--- a/test/test_autograd.py
+++ b/test/test_autograd.py
@@ -1077,6 +1077,8 @@
     (Resize, (S * S, S), ((S, S, S),)),
     (Diag, (), ((S, S),), '2d'),
     (Diag, (), ((S,),), '1d'),
+    (Diag, (1,), ((S, S),), '2d_1'),
+    (Diag, (2,), ((S, S),), '2d_2'),
     (Tril, (), ((S, S),)),
     (Tril, (2,), ((S, S),), 'idx'),
     (Triu, (), ((S, S),)),

torch/autograd/_functions/linalg.py

diff --git a/torch/autograd/_functions/linalg.py b/torch/autograd/_functions/linalg.py
index 4b522be..b0090f4 100644
--- a/torch/autograd/_functions/linalg.py
+++ b/torch/autograd/_functions/linalg.py
@@ -10,10 +10,10 @@
         self.diagonal_idx = diagonal_idx
 
     def forward(self, input):
-        return input.diag()
+        return input.diag(self.diagonal_idx)
 
     def backward(self, grad_output):
-        return grad_output.diag()
+        return grad_output.diag(self.diagonal_idx)
 
 
 class Tril(Function):