torch.diag bug fix (#1251)

pytorch · Apr 13, 2017 · cb66e9c · cb66e9c
1 parent c852883
commit cb66e9c
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/test/test_autograd.py b/test/test_autograd.py
@@ -1077,6 +1077,8 @@ def prod_single_zero(dim_size):
     (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),)),

diff --git a/torch/autograd/_functions/linalg.py b/torch/autograd/_functions/linalg.py
@@ -10,10 +10,10 @@ def __init__(self, diagonal_idx=0):
         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):