Q:

Finding column index of a PyTorch tensor that has most 1's

I have a PyTorch tensor a shaped like below:
import torch
a = torch.tensor([[[1., 0., 0., 0.]],
        [[0., 1., 0., 0.]],
        [[1., 0., 0., 0.]],
        [[0., 0., 0., 1.]],
        [[1., 0., 0., 0.]],
        [[0., 0., 0., 1.]],
        [[1., 0., 0., 0.]]])

Each row of the tensor a has 4 elements, 1's and 0's. Say I index the row and the column of this tensor accordingly. So for instance, the entry in the row 0 (uppermost row) is [[1., 0., 0., 0.]], whereas the entry in the column 3 (rightmost column) is [[0., 0., 0., 1., 0., 1., 0.]].
From a given tensor, I want to identify the index of the column where 1. appear most frequently. For example, for the tensor a, the index of such a column would be 0. If there are ties in the number of 1.'s, I still would like to get all of those tied column indices.
How can I do this task easily on Python?
Thank you,

A:

If your matrix only contains 0 and 1, you can sum the elements of each column and then search for the sum that is the largest:
import numpy as np

% sum over columns
sumsi = torch.sum(a, dim=1)

% find where maximum
col_idx = np.where(sumsi==np.max(sumsi))

