Hi All,
I was wondering if it’s at all possible to take the trace of matrix for a batch of matrices? For example, let’s say I have some Tensor of shape [B, N, N] and wish to find the trace along B for each [N, N] matrix?
If I use the torch.trace command, I get an error saying that it expected a matrix. I understand I could just iterate over dim=0, and take the trace of each matrix then add them to a vector. But is there a way to do this with one command? Which takes [B, N, N] and returns [B]?
Thanks in advance!