RuntimeError: Cost Function returns nan values in its 1th output
I am working on a temporal model to predict future events. Here is the link to my colab notebook. I am facing an issue trying to train the model. I am getting NaN values for Train and valid loss. The loss function is a joint loss consisting of the cross entropy loss and the squared loss.  Link to the blog here.
Tried following solution which didn`t worked -
Smaller learning rate - 0.01, 0.001, 0.0001
class cost_function():
    def __init__(self, yhat, y, L_2=0.001, logEps=1e-8):
        # logEps : log epsilon, very small positive value greater that 0.0
        # CE = - [ ln(y*)(y) + ln(1-y*)(1-y) ],

        self.yhat = yhat
        self.y = y
       
        self.logEps = logEps
        self.L_2 = L_2
        
        self.W_out = nn.Parameter(torch.randn(hiddenDimSize, numClass)*0.01)
        
    def cross_entropy(self):
        ce = -(self.y * torch.log(self.yhat + self.logEps) + (1. - self.y) * torch.log(1. - self.yhat + self.logEps))
        print(&quot;Inside CrossEntrophy Loss fn : &quot;, ce)
        return ce

    def prediction_loss(self):
        # return  (torch.sum(torch.sum(self.cross_entropy(), dim=0),dim=1)).float()/  lengths.float()
        
        tmp_tensor = torch.sum(self.cross_entropy(), dim=0)
        print(&quot;Inside PredictionLoss fn : Sum Dim 0&quot;, tmp_tensor)
        print(&quot;Inside PredictionLoss fn : Sum Dim 1&quot;, torch.sum(tmp_tensor,dim=1))
        print(&quot;Inside PredictionLoss fn : Final Result &quot;, (torch.sum(tmp_tensor,dim=1)).float()/  lengths.float())
        return (torch.sum(tmp_tensor,dim=1)).float()/  lengths.float()
        
    def cost(self):
        print(&quot;Inside Cost fn :&quot;, torch.mean(self.prediction_loss()) + self.L_2 * (self.W_out ** 2).sum())
        return torch.mean(self.prediction_loss()) + self.L_2 * (self.W_out ** 2).sum() # regularize
    

build_EHRNN class -  I have made modification in the forward method params to resolve the 'h' is not defined error.
torch.manual_seed(1)

class build_EHRNN(nn.Module):
    def __init__(self, inputDimSize=4894, hiddenDimSize=[200,200], batchSize=100, embSize=200, numClass=4894, dropout=0.5, logEps=1e-8):
        super(build_EHRNN, self).__init__()
        
        self.inputDimSize = inputDimSize
        self.hiddenDimSize = hiddenDimSize
        self.numClass = numClass
        self.embSize = embSize
        self.batchSize = batchSize
        self.dropout = nn.Dropout(p=0.5)
        self.logEps = logEps
        
        
        # Embedding inputs
        self.W_emb = nn.Parameter(torch.randn(self.inputDimSize, self.embSize).cuda())
        self.b_emb = nn.Parameter(torch.zeros(self.embSize).cuda())
        
        self.W_out = nn.Parameter(torch.randn(self.hiddenDimSize, self.numClass).cuda())
        self.b_out = nn.Parameter(torch.zeros(self.numClass).cuda())
         
        self.params = [self.W_emb, self.W_out, 
                       self.b_emb, self.b_out] 
    
    # def forward(self,x, y, h, lengths, mask):
    def forward(self,x, y, lengths, mask):
        self.emb = torch.tanh(torch.matmul(x, self.W_emb) + self.b_emb)
        input_values = self.emb
        self.outputs = [input_values]
        for i, hiddenSize in enumerate([self.hiddenDimSize, self.hiddenDimSize]):  # iterate over layers
            rnn = EHRNN(self.inputDimSize,hiddenSize,self.embSize,self.batchSize,self.numClass) # calculate hidden states
            hidden_state = []
            h = self.init_hidden().cuda()
            for i,seq in enumerate(input_values): # loop over sequences in each batch
                h = rnn(seq, h)                    
                hidden_state.append(h)    
            hidden_state = self.dropout(torch.stack(hidden_state))    # apply dropout between layers
            input_values = hidden_state
       
        y_linear = torch.matmul(hidden_state, self.W_out)  + self.b_out # fully connected layer
        yhat = F.softmax(y_linear, dim=1)  # yhat
        yhat = yhat*mask[:,:,None]   # apply mask
        
        # Loss calculation
        cross_entropy = -(y * torch.log(yhat + self.logEps) + (1. - y) * torch.log(1. - yhat + self.logEps))
        last_step = -torch.mean(y[-1] * torch.log(yhat[-1] + self.logEps) + (1. - y[-1]) * torch.log(1. - yhat[-1] + self.logEps))
        prediction_loss = torch.sum(torch.sum(cross_entropy, dim=0),dim=1)/ torch.cuda.FloatTensor(lengths)
        cost = torch.mean(prediction_loss) + 0.000001 * (self.W_out ** 2).sum() # regularize
        return (yhat, hidden_state, cost)

    def init_hidden(self):
        return torch.zeros(self.batchSize, self.hiddenDimSize)  # initial state

model training
artificalData_seqs = np.array(pickle.load(open(os.path.join(GOOGLE_DRV_PATH,BASE_DIR,'data.encodedDxs'),'rb')))
train, test, valid = load_data(artificalData_seqs, artificalData_seqs)

batchSize = 50     # decreased from 100 to 50
n_batches = int(np.ceil(float(len(train[0])) / float(batchSize)))-1
n_batches_valid = int(np.ceil(float(len(valid[0])) / float(batchSize)))-1
model = build_EHRNN(inputDimSize=4894, hiddenDimSize=200, batchSize=50, embSize=200, numClass=4894, dropout=0.5, logEps=1e-8)
model = model.to(device)



 import torch.nn.functional as F
import pdb

optimizer = torch.optim.Adadelta(model.parameters(), lr = 0.001, rho=0.95)
epochs = 10

counter = 0
# with torch.autograd.detect_anomaly():
for e in range(epochs):
    for x, y in train_dl:
        x, y , mask, lengths = padding(x, y, inputDimSize, numClass)
        output, h = model(x, mask)
        
        loss = cost_function(output, y).cost()
        # pdb.set_trace()
        loss.backward()
        print(&quot;loss &quot;,loss)
        nn.utils.clip_grad_norm_(model.parameters(), 5) # Constraining the weight matrix directly == regularization. 
        optimizer.step()
        optimizer.zero_grad()
    
    with torch.no_grad():
            model.eval()
            val_loss = []
            for x_valid, y_valid in valid_dl:
                    x_val, y_val, mask, lengths = padding(x_valid, y_valid, inputDimSize, numClass)
                    outputs_val, hidden_val = model(x_val,  mask)
                    loss = cost_function(outputs_val, y_val).cost()
                    val_loss.append(loss.item())
            model.train()

            print(&quot;Epoch: {}/{}...&quot;.format(e+1, epochs),
                                &quot;Step: {}...&quot;.format(counter),
                                &quot;Training Loss: {:.4f}...&quot;.format(loss.item()),
                                &quot;Val Loss: {:.4f}&quot;.format(torch.mean(torch.tensor(val_loss))))

Error (Start getting NaNs in the loss)
Inside PredictionLoss fn : Sum Dim 0 tensor([[0.1008, 0.1539, 0.1211,  ..., 0.1533, 0.1218, 0.1418],
        [0.0253, 0.0449, 0.0249,  ..., 0.0439, 0.0134, 0.0332],
        [0.0306, 0.0799, 0.0570,  ..., 0.0790, 0.0484, 0.0678],
        ...,
        [0.0253, 0.0450, 0.0249,  ..., 0.0439, 0.0134, 0.0332],
        [0.0253, 0.0450, 0.0249,  ..., 0.0439, 0.0134, 0.0332],
        [0.0038, 0.0106, 0.0106,  ..., 0.0098, 0.0004, 0.0106]],
       grad_fn=&lt;SumBackward1&gt;)
Inside PredictionLoss fn : Sum Dim 1 tensor([  372.4754,   133.2620,   219.1195,    37.5425,   141.3354,    37.5070,
          229.2947,     0.0000,   379.1829,   217.3962,    80.1226,    37.5074,
          138.4665,    82.1034,    89.7893,    81.8173,    92.8159,   141.8856,
           95.9898,   216.0511,   133.2535,   385.0391,   369.4958,   244.9362,
           37.5088,    37.5087,   141.6083,     0.0000,    95.3367,    37.5074,
          735.0569,   378.0407,    37.5135,    40.7778,    82.0872,   225.9998,
          216.6189,   379.0732,    81.4742,   144.4226,    93.3905,   214.0228,
           37.5078,   224.0793,    88.3753,    41.2919,   140.4855,    37.5086,
          226.6366,   148.7171,   137.9226, 13887.5811,    81.1428,    84.6804,
          226.6779,    37.5065,   223.8841,   220.5979,    83.2484,    37.5080,
           84.5247,   384.2115,    80.1173,     0.0000,   146.9714,    37.6982,
          134.6618,    84.1838,    37.5421,   730.5516,    37.5085,     0.0000,
          215.1523,   136.5673,    81.2887,    37.5080,    94.4181,   140.6268,
          133.9295,   136.2485,    80.1173,   386.2103,    39.0282,     0.0000,
           37.5055,    42.1506,    80.1662,   228.5819,    39.3403,   138.7672,
         1768.6033,   143.5350,    40.2060,   147.7809,    37.5085,   380.9214,
          750.6883,   141.0447,   136.9028,    37.5049],
       grad_fn=&lt;SumBackward1&gt;)
Inside PredictionLoss fn : lengths tensor([5., 3., 4., 1., 3., 1., 4., 0., 5., 4., 2., 1., 3., 2., 2., 2., 2., 3.,
        2., 4., 3., 5., 5., 4., 1., 1., 3., 0., 2., 1., 6., 5., 1., 1., 2., 4.,
        4., 5., 2., 3., 2., 4., 1., 4., 2., 1., 3., 1., 4., 3., 3., 9., 2., 2.,
        4., 1., 4., 4., 2., 1., 2., 5., 2., 0., 3., 1., 3., 2., 1., 6., 1., 0.,
        4., 3., 2., 1., 2., 3., 3., 3., 2., 5., 1., 0., 1., 1., 2., 4., 1., 3.,
        7., 3., 1., 3., 1., 5., 6., 3., 3., 1.])
Inside PredictionLoss fn : Final Result  tensor([  74.4951,   44.4207,   54.7799,   37.5425,   47.1118,   37.5070,
          57.3237,       nan,   75.8366,   54.3491,   40.0613,   37.5074,
          46.1555,   41.0517,   44.8946,   40.9086,   46.4080,   47.2952,
          47.9949,   54.0128,   44.4178,   77.0078,   73.8992,   61.2340,
          37.5088,   37.5087,   47.2028,       nan,   47.6683,   37.5074,
         122.5095,   75.6081,   37.5135,   40.7778,   41.0436,   56.5000,
          54.1547,   75.8146,   40.7371,   48.1409,   46.6952,   53.5057,
          37.5078,   56.0198,   44.1876,   41.2919,   46.8285,   37.5086,
          56.6591,   49.5724,   45.9742, 1543.0646,   40.5714,   42.3402,
          56.6695,   37.5065,   55.9710,   55.1495,   41.6242,   37.5080,
          42.2623,   76.8423,   40.0586,       nan,   48.9905,   37.6982,
          44.8873,   42.0919,   37.5421,  121.7586,   37.5085,       nan,
          53.7881,   45.5224,   40.6443,   37.5080,   47.2090,   46.8756,
          44.6432,   45.4162,   40.0587,   77.2421,   39.0282,       nan,
          37.5055,   42.1506,   40.0831,   57.1455,   39.3403,   46.2557,
         252.6576,   47.8450,   40.2060,   49.2603,   37.5085,   76.1843,
         125.1147,   47.0149,   45.6343,   37.5049], grad_fn=&lt;DivBackward0&gt;)
Inside PredictionLoss fn : Sum Dim 0 tensor([[nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        ...,
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan]], grad_fn=&lt;SumBackward1&gt;)
Inside PredictionLoss fn : Sum Dim 1 tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan], grad_fn=&lt;SumBackward1&gt;)
Inside PredictionLoss fn : lengths tensor([2., 1., 1., 2., 3., 5., 3., 3., 1., 4., 2., 2., 2., 3., 2., 2., 5., 2.,
        4., 2., 1., 4., 2., 5., 3., 0., 4., 7., 6., 4., 4., 1., 7., 1., 3., 3.,
        5., 3., 5., 5., 4., 2., 2., 4., 1., 5., 6., 2., 5., 5., 2., 1., 3., 1.,
        4., 4., 3., 3., 0., 2., 2., 4., 2., 2., 1., 0., 3., 2., 3., 0., 1., 2.,
        4., 4., 5., 1., 1., 3., 3., 2., 5., 0., 2., 6., 5., 5., 5., 1., 2., 4.,
        0., 2., 6., 1., 2., 0., 1., 1., 2., 3.])
Inside PredictionLoss fn : Final Result  tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan], grad_fn=&lt;DivBackward0&gt;)
Inside PredictionLoss fn : Sum Dim 0 tensor([[nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        ...,
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan]], grad_fn=&lt;SumBackward1&gt;)
Inside PredictionLoss fn : Sum Dim 1 tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan], grad_fn=&lt;SumBackward1&gt;)
Inside PredictionLoss fn : lengths tensor([3., 2., 4., 4., 6., 1., 1., 4., 2., 2., 4., 2., 1., 1., 1., 1., 3., 1.,
        2., 3., 4., 2., 2., 3., 2., 0., 1., 3., 4., 1., 2., 3., 4., 3., 3., 1.,
        2., 3., 3., 2., 1., 4., 6., 3., 4., 2., 3., 0., 1., 1., 3., 7., 2., 3.,
        4., 2., 3., 0., 3., 3., 2., 1., 1., 3., 6., 2., 2., 3., 2., 3., 3., 2.,
        1., 2., 5., 3., 4., 2., 5., 3., 5., 5., 5., 3., 2., 1., 4., 3., 3., 3.,
        1., 4., 3., 2., 5., 3., 6., 4., 3., 2.])
Inside PredictionLoss fn : Final Result  tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan], grad_fn=&lt;DivBackward0&gt;)
Inside PredictionLoss fn : Sum Dim 0 tensor([[nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        ...,
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan]], grad_fn=&lt;SumBackward1&gt;)
Inside PredictionLoss fn : Sum Dim 1 tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan], grad_fn=&lt;SumBackward1&gt;)
Inside PredictionLoss fn : lengths tensor([4., 6., 2., 2., 4., 5., 0., 3., 2., 3., 4., 3., 3., 3., 4., 3., 0., 1.,
        2., 2., 4., 2., 4., 4., 1., 0., 3., 0., 5., 5., 1., 2., 3., 2., 2., 3.,
        1., 2., 1., 2., 3., 1., 3., 1., 3., 0., 4., 2., 1., 6., 2., 2., 0., 1.,
        4., 4., 2., 2., 0., 3., 5., 5., 5., 5., 3., 4., 4., 4., 5., 4., 2., 2.,
        3., 1., 5., 6., 3., 4., 1., 3., 5., 1., 3., 2., 3., 6., 2., 4., 3., 2.,
        3., 2., 4., 1., 2., 2., 1., 2., 3., 4.])
Inside PredictionLoss fn : Final Result  tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan], grad_fn=&lt;DivBackward0&gt;)
Inside PredictionLoss fn : Sum Dim 0 tensor([[nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        ...,
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan]], grad_fn=&lt;SumBackward1&gt;)
Inside PredictionLoss fn : Sum Dim 1 tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan], grad_fn=&lt;SumBackward1&gt;)
Inside PredictionLoss fn : lengths tensor([1., 5., 1., 5., 3., 0., 1., 4., 3., 3., 5., 5., 2., 0., 1., 2., 3., 2.,
        2., 2., 2., 2., 2., 2., 2., 3., 4., 1., 4., 2., 1., 3., 4., 0., 4., 2.,
        0., 4., 1., 1., 2., 3., 3., 5., 1., 3., 2., 3., 1., 1., 3., 6., 3., 0.,
        2., 3., 2., 3., 3., 4., 1., 5., 2., 6., 2., 1., 5., 3., 2., 1., 3., 5.,
        2., 0., 3., 0., 1., 3., 3., 2., 3., 1., 2., 4., 1., 2., 7., 2., 7., 2.,
        2., 2., 1., 1., 1., 2., 4., 2., 2., 3.])
Inside PredictionLoss fn : Final Result  tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan], grad_fn=&lt;DivBackward0&gt;)
Inside PredictionLoss fn : Sum Dim 0 tensor([[nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        ...,
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan]], grad_fn=&lt;SumBackward1&gt;)
Inside PredictionLoss fn : Sum Dim 1 tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan], grad_fn=&lt;SumBackward1&gt;)
Inside PredictionLoss fn : lengths tensor([2., 2., 1., 4., 3., 2., 2., 4., 5., 2., 3., 1., 1., 2., 3., 1., 3., 1.,
        3., 1., 6., 4., 2., 2., 0., 4., 1., 0., 2., 0., 1., 1., 1., 1., 2., 2.,
        3., 1., 3., 2., 0., 1., 3., 2., 4., 1., 2., 1., 6., 2., 1., 2., 1., 2.,
        2., 1., 1., 1., 4., 1., 0., 2., 2., 2., 3., 6., 5., 1., 1., 3., 3., 3.,
        1., 4., 4., 2., 4., 2., 3., 7., 1., 1., 2., 3., 0., 0., 1., 4., 2., 3.,
        1., 2., 3., 2., 7., 0., 2., 3., 4., 4.])
Inside PredictionLoss fn : Final Result  tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan], grad_fn=&lt;DivBackward0&gt;)
Inside PredictionLoss fn : Sum Dim 0 tensor([[nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        ...,
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan]], grad_fn=&lt;SumBackward1&gt;)
Inside PredictionLoss fn : Sum Dim 1 tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan], grad_fn=&lt;SumBackward1&gt;)
Inside PredictionLoss fn : lengths tensor([1., 2., 6., 7., 2., 1., 1., 3., 1., 2., 2., 4., 3., 3., 4., 3., 2., 2.,
        5., 2., 1., 3., 2., 2., 2., 1., 4., 3., 2., 2., 1., 4., 2., 0., 2., 2.,
        1., 3., 4., 4., 5., 1., 5., 6., 1., 3., 1., 3., 0., 1., 6., 1., 2., 0.,
        3., 2., 2., 3., 4., 3., 3., 2., 3., 2., 2., 3., 4., 3., 2., 0., 3., 1.,
        1., 3., 1., 4., 7., 2., 2., 1., 1., 1., 2., 3., 3., 2., 2., 3., 1., 3.,
        3., 1., 1., 3., 1., 2., 7., 3., 1., 6.])
Inside PredictionLoss fn : Final Result  tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan], grad_fn=&lt;DivBackward0&gt;)
Inside PredictionLoss fn : Sum Dim 0 tensor([[nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        ...,
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan]], grad_fn=&lt;SumBackward1&gt;)
Inside PredictionLoss fn : Sum Dim 1 tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan], grad_fn=&lt;SumBackward1&gt;)
Inside PredictionLoss fn : lengths tensor([2., 2., 4., 2., 5., 1., 2., 4., 2., 2., 2., 5., 1., 2., 1., 4., 1., 2.,
        4., 3., 4., 0., 2., 5., 2., 6., 2., 6., 3., 4., 3., 5., 7., 4., 3., 5.,
        1., 3., 2., 3., 2., 2., 2., 3., 3., 3., 2., 3., 1., 4., 3., 3., 2., 1.,
        6., 9., 2., 1., 6., 3., 1., 5., 1., 6., 2., 2., 6., 2., 4., 2., 4., 3.,
        2., 2., 4., 1., 2., 2., 2., 1., 3., 1., 3., 2., 3., 3., 2., 2., 4., 3.,
        3., 3., 5., 4., 3., 2., 4., 2., 1., 8.])
Inside PredictionLoss fn : Final Result  tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan], grad_fn=&lt;DivBackward0&gt;)
Inside PredictionLoss fn : Sum Dim 0 tensor([[nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        ...,
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan]], grad_fn=&lt;SumBackward1&gt;)
Inside PredictionLoss fn : Sum Dim 1 tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan], grad_fn=&lt;SumBackward1&gt;)
Inside PredictionLoss fn : lengths tensor([1., 3., 2., 3., 3., 2., 7., 4., 3., 1., 1., 5., 1., 3., 4., 4., 1., 5.,
        1., 3., 3., 1., 4., 3., 0., 4., 4., 2., 4., 1., 2., 4., 3., 2., 3., 2.,
        2., 2., 2., 3., 3., 2., 5., 1., 2., 1., 2., 2., 3., 2., 2., 2., 3., 3.,
        3., 5., 1., 4., 8., 4., 0., 2., 4., 2., 0., 1., 4., 4., 1., 5., 0., 1.,
        3., 1., 2., 3., 3., 4., 3., 4., 2., 4., 2., 3., 4., 1., 1., 2., 4., 1.,
        1., 0., 0., 0., 4., 3., 1., 3., 4., 3.])
Inside PredictionLoss fn : Final Result  tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan,
        nan, nan, nan, nan], grad_fn=&lt;DivBackward0&gt;)
Inside PredictionLoss fn : Sum Dim 0 tensor([[nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        ...,
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan],
        [nan, nan, nan,  ..., nan, nan, nan]], grad_fn=&lt;SumBackward1&gt;)


The problem is the values in lengths variable.
In your cost_function.prediction_loss you divide the cross entropy loss by the length of each sequece: (torch.sum(tmp_tensor,dim=1)).float()/  lengths.float().
However, if you look at the values of your lengths tensor:

Inside PredictionLoss fn : lengths tensor([5., 3., 4., 1., 3., 1., 4., 0., 5., 4., 2., 1., 3., 2., 2., 2., 2., 3.,
   2., 4., 3., 5., 5., 4., 1., 1., 3., 0., 2., 1., 6., 5., 1., 1., 2., 4.,
   4., 5., 2., 3., 2., 4., 1., 4., 2., 1., 3., 1., 4., 3., 3., 9., 2., 2.,
   4., 1., 4., 4., 2., 1., 2., 5., 2., 0., 3., 1., 3., 2., 1., 6., 1., 0.,
   4., 3., 2., 1., 2., 3., 3., 3., 2., 5., 1., 0., 1., 1., 2., 4., 1., 3.,
   7., 3., 1., 3., 1., 5., 6., 3., 3., 1.])


You will notice that some of the entries are 0(!). The corresponding values in the loss function are also zero (no loss for zero-length sequence). When you divide zero by zero you get nan.

Some good practices for coding:

If possible, use library functions instead of re-implementing things. These functions are usually tested and optimized and are more numerically stable.
For instance, you can use torch.nn.CrossEntropyLoss that combines both the cross entropy loss and the softmax in a numerically robust manner.

the variable lengths used for loss computation is not explicitly an argument of the loss function or a class member. You should make it an explicit argument.


