Issue
Consider I have 2D Tensor, index_in_batch * diag_ele
.
How can I get a 3D Tensor index_in_batch * Matrix
(who is a diagonal matrix, construct by drag_ele)?
The torch.diag()
construct diagonal matrix only when input is 1D, and return diagonal element when input is 2D.
Solution
import torch
a = torch.rand(2, 3)
print(a)
b = torch.eye(a.size(1))
c = a.unsqueeze(2).expand(*a.size(), a.size(1))
d = c * b
print(d)
Output
0.5938 0.5769 0.0555
0.9629 0.5343 0.2576
[torch.FloatTensor of size 2x3]
(0 ,.,.) =
0.5938 0.0000 0.0000
0.0000 0.5769 0.0000
0.0000 0.0000 0.0555
(1 ,.,.) =
0.9629 0.0000 0.0000
0.0000 0.5343 0.0000
0.0000 0.0000 0.2576
[torch.FloatTensor of size 2x3x3]
Answered By - Wasi Ahmad
0 comments:
Post a Comment
Note: Only a member of this blog may post a comment.