title: How to get the class with highest confidence from a tensor?
body: <p>I want to obtain the class with the highest confidence.
  Heres the code that performs classification task:</p>
  <pre><code>names = ['class A', 'class B', 'class C']
  
  def classify_face(image):
      device = torch.device(&quot;cpu&quot;)
      img = process_image(image)
      print('Image processed')
      # img = image.unsqueeze_(0)
      # img = image.float()
      
      pred = model(img)[0]
      
      # Apply NMS
      pred = non_max_suppression(pred, 0.4, 0.5, classes = [0, 1, 2], agnostic = None )
      if classify:
          pred = apply_classifier(pred, modelc, img, im0s)
      #print(pred)
      
      model.eval()
      model.cpu()
      
      print(pred)
          
      # output = non_max_suppression(output, 0.4, 0.5, classes = class_names, agnostic = False)
          
      #_, predicted = torch.max(output[0], 1)
      #print(predicted.data[0], &quot;predicted&quot;)
  
      classification = torch.cat(pred)[:, -1]
      index = int(classification)
      print(names[index])
      return names[index]
  </code></pre>
  <p>The above code works perfectly when the pred has a 1D tensor. But I get an error if the tensor size is more than that.</p>
  <p>There are 5 elements: <code>x1</code>, <code>y1</code>, <code>x2</code>, <code>y2</code>, <code>confidence</code>, and <code>class</code>.</p>
  <p>For example:</p>
  <pre><code>pred = [torch.tensor([[212.38568, 117.47020, 339.35773, 266.00513, 0.74144, 2.00000],
                        [214.60651, 118.50694, 339.90192, 265.91696, 0.94277, 0.00000]])]
  </code></pre>
  <p>Error:</p>
  <pre><code>Traceback (most recent call last):
    File &quot;WEBCAM_DETECT.py&quot;, line 172, in &lt;module&gt;
      label = classify_face(frame)
    File &quot;WEBCAM_DETECT.py&quot;, line 154, in classify_face
      index = int(classification)
  ValueError: only one element tensors can be converted to Python scalars
  </code></pre>
  <p>So I want to access the class with highest confidence. Please let me know how to do that or if there is a better way to do.</p>