Argmax of 2d vector on C++
I am working on python/pytorch and I have an example like
2d vector a
 
    |
    v
  dim-0  ---&gt; -----&gt; dim-1 ------&gt; -----&gt; --------&gt; dim-1
    |   [[-1.7739,  0.8073,  0.0472, -0.4084],
    v    [ 0.6378,  0.6575, -1.2970, -0.0625],
    |    [ 1.7970, -1.3463,  0.9011, -0.8704],
    v    [ 1.5639,  0.7123,  0.0385,  1.8410]]
    |
    v

Then, the argmax with the index of 1 will be
# argmax (indices where max values are present) along dimension-1
In [215]: torch.argmax(a, dim=1)
Out[215]: tensor([1, 1, 0, 3])

My question is that given the 2d vector a as above, how could I implement argmax function on C++ to give me same output as above? Thanks for reading
This is what I did
   

    vector&lt;vector&lt;float&gt;&gt; a_vect

    {
      {-1.7739,  0.8073,  0.0472, -0.4084},
      {0.6378,  0.6575, -1.2970, -0.0625},
      {1.7970, -1.3463,  0.9011, -0.8704},
      {1.5639,  0.7123,  0.0385,  1.8410}
     };
     std::vector&lt;int&gt;::iterator max = max_element(a_vect.begin() , a_vect.end()-a_vect.begin());


You can use std::max_element to find the index in each sub vector
#include &lt;algorithm&gt;
#include &lt;iostream&gt;
#include &lt;vector&gt;
using std::vector;

int main()
{
  vector&lt;vector&lt;float&gt;&gt; a_vect=
  {
    {-1.7739,  0.8073,  0.0472, -0.4084},
    {0.6378,  0.6575, -1.2970, -0.0625},
    {1.7970, -1.3463,  0.9011, -0.8704},
    {1.5639,  0.7123,  0.0385,  1.8410}
  };
  
  vector&lt;int&gt; max_index;
  for(auto&amp; v:a_vect)
    max_index.push_back(std::max_element(v.begin(),v.end())-v.begin());
  
  for(auto i:max_index)
    std::cout &lt;&lt; i &lt;&lt; ' '; // 1 1 0 3 
}

