title: Custom cross-entropy loss in pytorch
body: <p>I have done a custom implementation of the pytorch cross-entropy loss function (as I need more flexibility to be introduced later). The model I intend to train with this will need a considerable amount of time to train and the resources available can't be used to merely test if the function is correct implementation. I have implemented vectorized implementation as it will be quicker to run.</p>
  
  <p>Following is my code for the same:</p>
  
  <pre class="lang-py prettyprint-override"><code>def custom_cross(my_pred,true,batch_size=BATCH_SIZE):
      loss= -torch.mean(torch.sum(true.view(batch_size, -1) * torch.log(my_pred.view(batch_size, -1)), dim=1))
      return loss
  </code></pre>
  
  <p>I will really appreciate if you can suggest a more optimized implementation of the same or if I am making a mistake in the present one. The model will use a Nvidia Tesla K-80 to train.</p>