title: PyTorch softmax with dim
body: <p>Which dimension should softmax be applied to ?</p>
  
  <p>This code : </p>
  
  <pre><code>%reset -f
  
  import torch.nn as nn
  import numpy as np
  import torch
  
  my_softmax = nn.Softmax(dim=-1)
  
  mu, sigma = 0, 0.1 # mean and standard deviation
  
  train_dataset = []
  image = []
  image_x = np.random.normal(mu, sigma, 24).reshape((3 , 4, 2))
  train_dataset.append(image_x)
  
  x = torch.tensor(train_dataset).float()
  
  print(x)
  print(my_softmax(x))
  my_softmax = nn.Softmax(dim=1)
  print(my_softmax(x))
  </code></pre>
  
  <p>prints following : </p>
  
  <pre><code>tensor([[[[-0.1500,  0.0243],
            [ 0.0226,  0.0772],
            [-0.0180, -0.0278],
            [ 0.0782, -0.0853]],
  
           [[-0.0134, -0.1139],
            [ 0.0385, -0.1367],
            [-0.0447,  0.1493],
            [-0.0633, -0.2964]],
  
           [[ 0.0123,  0.0061],
            [ 0.1086, -0.0049],
            [-0.0918, -0.1308],
            [-0.0100,  0.1730]]]])
  tensor([[[[ 0.4565,  0.5435],
            [ 0.4864,  0.5136],
            [ 0.5025,  0.4975],
            [ 0.5408,  0.4592]],
  
           [[ 0.5251,  0.4749],
            [ 0.5437,  0.4563],
            [ 0.4517,  0.5483],
            [ 0.5580,  0.4420]],
  
           [[ 0.5016,  0.4984],
            [ 0.5284,  0.4716],
            [ 0.5098,  0.4902],
            [ 0.4544,  0.5456]]]])
  tensor([[[[ 0.3010,  0.3505],
            [ 0.3220,  0.3665],
            [ 0.3445,  0.3230],
            [ 0.3592,  0.3221]],
  
           [[ 0.3450,  0.3053],
            [ 0.3271,  0.2959],
            [ 0.3355,  0.3856],
            [ 0.3118,  0.2608]],
  
           [[ 0.3540,  0.3442],
            [ 0.3509,  0.3376],
            [ 0.3200,  0.2914],
            [ 0.3289,  0.4171]]]])
  </code></pre>
  
  <p>So first tensor is prior to softmax being applied, second tensor is result of softmax applied to tensor with dim=-1 and third tensor is result of softmax applied to tensor with dim=1 .</p>
  
  <p>For result of first softmax can see corresponding elements sum to 1, for example [ 0.4565,  0.5435] -> 0.4565 + 0.5435 == 1.</p>
  
  <p>What is summing to 1 as result of of second softmax ?</p>
  
  <p>Which dim value should I choose ?</p>
  
  <p>Update : The dimension <code>(3 , 4, 2)</code> corresponds to image dimension where 3 is the RGB value , 4 is the number of horizontal pixels (width) , 2 is the number of vertical pixels (height). This is an image classification problem. I'm using cross entropy loss function. Also, I'm using softmax in final layer in order to back-propagate probabilities.</p>