title: PyTorch custom loss function
body: <p>How should a custom loss function be implemented ? Using below code is causing error :</p>
  <pre><code>import torch
  import torch.nn as nn
  import torchvision
  import torchvision.transforms as transforms
  import numpy as np
  import matplotlib.pyplot as plt
  import torch.utils.data as data_utils
  import torch.nn as nn
  import torch.nn.functional as F
  
  num_epochs = 20
  
  x1 = np.array([0,0])
  x2 = np.array([0,1])
  x3 = np.array([1,0])
  x4 = np.array([1,1])
  
  num_epochs = 200
  
  class cus2(torch.nn.Module):
      
      def __init__(self):
          super(cus2,self).__init__()
      
      def forward(self, outputs, labels):
          # reshape labels to give a flat vector of length batch_size*seq_len
          labels = labels.view(-1)  
  
          # mask out 'PAD' tokens
          mask = (labels &gt;= 0).float()
  
          # the number of tokens is the sum of elements in mask
          num_tokens = int(torch.sum(mask).data[0])
  
          # pick the values corresponding to labels and multiply by mask
          outputs = outputs[range(outputs.shape[0]), labels]*mask
  
          # cross entropy loss for all non 'PAD' tokens
          return -torch.sum(outputs)/num_tokens
  
  
  x = torch.tensor([x1,x2,x3,x4]).float()
  
  y = torch.tensor([0,1,1,0]).long()
  
  train = data_utils.TensorDataset(x,y)
  train_loader = data_utils.DataLoader(train , batch_size=2 , shuffle=True)
  
  device = 'cpu'
  
  input_size = 2
  hidden_size = 100 
  num_classes = 2
  
  learning_rate = .0001
  
  class NeuralNet(nn.Module) : 
      def __init__(self, input_size, hidden_size, num_classes) : 
          super(NeuralNet, self).__init__()
          self.fc1 = nn.Linear(input_size , hidden_size)
          self.relu = nn.ReLU()
          self.fc2 = nn.Linear(hidden_size , num_classes)
  
      def forward(self, x) : 
          out = self.fc1(x)
          out = self.relu(out)
          out = self.fc2(out)
          return out
          
  for i in range(0 , 1) :
          
          model = NeuralNet(input_size, hidden_size, num_classes).to(device)
          
          criterion = nn.CrossEntropyLoss()
  #         criterion = Regress_Loss()
  #         criterion = cus2()
          optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
          
          total_step = len(train_loader)
          for epoch in range(num_epochs) : 
              for i,(images , labels) in enumerate(train_loader) : 
                  images = images.reshape(-1 , 2).to(device)
                  labels = labels.to(device)
                  
                  outputs = model(images)
                  loss = criterion(outputs , labels)
                  
                  optimizer.zero_grad()
                  loss.backward()
                  optimizer.step()
  #                 print(loss)
                  
          outputs = model(x)
          
          print(outputs.data.max(1)[1])
  </code></pre>
  <p>makes perfect predictions on training data :</p>
  <pre><code>tensor([0, 1, 1, 0])
  </code></pre>
  <p>Using a custom loss function from <a href="https://web.archive.org/web/20181129040718/https://cs230-stanford.github.io/pytorch-nlp.html" rel="nofollow noreferrer">here</a>:</p>
  <p><a href="https://i.stack.imgur.com/Eh3bZ.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/Eh3bZ.png" alt="image of the code used for the cus2 class" /></a></p>
  <p>is implemented in above code as <code>cus2</code></p>
  <p>Un-commenting code <code>#         criterion = cus2()</code> to use this loss function returns :</p>
  <pre><code>tensor([0, 0, 0, 0])
  </code></pre>
  <p>A warning is also returned :</p>
  <blockquote>
  <p>UserWarning: invalid index of a 0-dim tensor. This will be an error in
  PyTorch 0.5. Use tensor.item() to convert a 0-dim tensor to a Python
  number</p>
  </blockquote>
  <p>I've not implemented the custom loss function correctly ?</p>