fix for 1x1 inputs again
@@ -20,7 +20,7 @@ class AddDiag(Function):
if self.needs_input_grad[1]:
diag_grad = grad_output.new().resize_(1)
if grad_output.numel() == 1:
- diag_grad.fill_(grad_output[0])
+ diag_grad.fill_(grad_output.squeeze()[0])
else:
diag_grad.fill_(grad_output.trace())
