1. 程式人生 > >pytorch常用函式組合

pytorch常用函式組合

一、Find the indexes of some elements in a tenosr that fulfill some condition

Example

tensor:A=[[1,2,3],[-1,-2,-3],[0,0,0]]

condition:value of element in A > 0

output:indexes=[[0,0],[0,1],[0,2]]

indexes_list=torch.nonzero(torch.where(A>0,torch.ones(A.size()),torch.zeros(A.size())))