I am training a model with conv1d on top of the tdnn layers, but when i see the values in conv_tdnn in TDNNbase forward fxn after the first batch is executed, weights seem fine. but from second batch, When I checked the kernels/weights which I created and registered as parameters, the weights actually become NaN.  Actually for the first batch it works fine but after the optimization step i.e from second batch the weights (self.kernel and self.bias) becomes “Nan”
How do we fix this in pytorch
class TDNNBase(nn.Module):
    def __init__(self, in_channels, out_channels, context, full_context=True):     
        self.count = 0
        super(TDNNBase, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        # Check valid Context
        #self.context = context
        self.full_context = full_context
        self.is_context_valid(context)       
      
        """TDNN Layer  with Context = [-k,+k] or {-k, k}"""

        self.kernel_size, context = self.get_cont_ker_size(context, self.full_context)
        self.register_buffer("context", torch.LongTensor(context))
        self.kernel= nn.Parameter(torch.randn(self.out_channels, self.in_channels, self.kernel_size))
        self.bias  = nn.Parameter(torch.randn(self.out_channels)) 
        nn.init.xavier_uniform_(self.kernel)
        nn.init.xavier_uniform_(torch.unsqueeze(self.bias, 0))


        def forward(self, input):
          conv_tdnn = self.tdnnlayer(input, self.kernel, self.context, self.bias)
          return conv_tdnn

    """TDNN Convolution Layer"""
    def tdnn_layer(self, data_batch, kernel, context, bias):
        """
        Context wise Convolution: This is the time Delay Neural Network layer
        built on top of the 1D Convolution.Since Normal 1D Convolution is
        performed in the sequence, so in order to select the random context
        (i.e particular frames),we need to select only those particular Frames
        and perform the Convolution with only that particular Context (Only on those
        particular frames) """        

        input_shape = data_batch.size()        
        # Input Shape is batch size X Feature Dimension X Sequence(Frames) Length
        [batch_size, feature_dim, input_seq_length] = input_shape   
        num_frames_to_convolve = self.get_num_frames(context, input_seq_length)
        x_conv_tdnn = Variable(torch.zeros(batch_size, kernel.size()[0], len(num_frames_to_convolve)))  
        x_conv_tdnn[torch.isnan(x_conv_tdnn)] = 0           #Set Nan = 0      
        for i, step in enumerate(num_frames_to_convolve):
            features_to_convolve_indices = torch.index_select(data_batch, 2, context + step)  #To selecet features with context window using 
            x_conv_tdnn[:, :, i] = F.conv1d(features_to_convolve_indices, kernel,bias = bias, padding=1)[:,:,0]#.mean(dim=2)            
        return x_conv_tdnn   

class TDNN(nn.Module):
    def __init__(self):
        super(TimeDelayNN, self).__init__()
        self.tdnn1 = TDNNBase(13,512,[-2,2],full_context= True)  #in_channels = 40 is Feature Dimensions of Size 120 and
        self.bn1   = nn.BatchNorm1d(512)
        self.tdnn2 = TDNNBase(512, 64, [-1, 2],full_context=False) #out_channels = N Feature Maps/Kernels
        self.bn2   = nn.BatchNorm1d(64)
        self.conv1 = nn.Conv1d(64,16,1, padding=0)
        self.fc1   = nn.Linear(16,5)


    def forward(self, input):
        input = self.tdnn1(input)
        input = F.relu(input)
        input = self.bn1(input)
        input = self.tdnn2(input)
        input = F.relu(input)
        input = self.bn2(input)        
        input = self.conv1d1(input)
        input = F.relu(input)        
        input = self.fc1(input)    

        return input

#Training 
def train_fxn(**kwargs):    
    for i, (batch, label, _) in enumerate(self.train_loader):
        model.train()
        batch = batch.to(device)
        label = label.to(device)
        optim.zero_grad()                    #Zero out graidnets
        output = model(batch)                # Forward Pass
        loss = criterion(output, label)      #Loss
        loss.backward()
        #torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
        optim.step()
        _, predictions = torch.max(output.data, 1)        

        total += label.size(0)
        correct += (predictions == label).sum().item()
        running_train_loss += loss.item() * batch.size(0)
            
    train_avg_loss = running_train_loss / len(self.train_loader)
    return train_avg_loss, correct, total



#Output 
**After first batch at conv_tdnn (at return of TDNNbase forward function)**
tensor([[[ 0.0886,  0.0941,  0.0760,  ..., -0.1054, -0.1054, -0.1054],
         [-0.2494, -0.2688, -0.2743,  ..., -0.0443, -0.0443, -0.0443],
         [ 0.0373, -0.0574, -0.1424,  ..., -0.0450, -0.0450, -0.0450],
         ...,
         [ 0.1292,  0.0439, -0.0055,  ...,  0.0344,  0.0344,  0.0344],
         [ 0.1009,  0.1313,  0.0728,  ...,  0.0845,  0.0845,  0.0845],
         [-0.0864, -0.0101, -0.0266,  ...,  0.0706,  0.0706,  0.0706]],

        [[-0.0520, -0.0331, -0.0263,  ..., -0.1054, -0.1054, -0.1054],
         [-0.0542, -0.0904, -0.1329,  ..., -0.0443, -0.0443, -0.0443],
         [-0.2399, -0.2217, -0.1750,  ..., -0.0450, -0.0450, -0.0450],
         ...,
         [ 0.1098,  0.1261,  0.1243,  ...,  0.0344,  0.0344,  0.0344],
         [-0.0078,  0.0465,  0.0430,  ...,  0.0845,  0.0845,  0.0845],
         [-0.0522, -0.0356, -0.0696,  ...,  0.0706,  0.0706,  0.0706]],

        [[-0.0368, -0.0417, -0.0544,  ..., -0.1054, -0.1054, -0.1054],
         [-0.1158, -0.1055, -0.1214,  ..., -0.0443, -0.0443, -0.0443],
         [-0.2064, -0.1804, -0.1543,  ..., -0.0450, -0.0450, -0.0450],
         ...,
         [ 0.0432,  0.0606,  0.0588,  ...,  0.0344,  0.0344,  0.0344],
         [-0.0737, -0.0635, -0.0451,  ...,  0.0845,  0.0845,  0.0845],
         [-0.1350, -0.1560, -0.1374,  ...,  0.0706,  0.0706,  0.0706]],

        ...,

        [[-0.0406, -0.0759, -0.0934,  ...,  0.0270,  0.0437, -0.0089],
         [-0.0683, -0.0960, -0.1011,  ..., -0.4773, -0.4974, -0.4225],
         [-0.1967, -0.1860, -0.2016,  ...,  0.0212,  0.1038,  0.0512],
         ...,
         [ 0.0874,  0.0988,  0.0706,  ...,  0.1736,  0.1758,  0.1879],
         [-0.0339, -0.0344,  0.0466,  ..., -0.0173, -0.0020,  0.0231],
         [-0.0484, -0.0529,  0.0136,  ..., -0.3070, -0.3332, -0.3312]],

        [[ 0.0814,  0.0658,  0.0750,  ..., -0.0087,  0.0375,  0.0046],
         [-0.2866, -0.2804, -0.2798,  ..., -0.1776, -0.0918, -0.1062],
         [-0.3017, -0.2633, -0.2088,  ..., -0.1002, -0.1149, -0.1084],
         ...,
         [ 0.0134,  0.0475,  0.0628,  ...,  0.1704,  0.1373,  0.1426],
         [-0.0346, -0.0488, -0.0529,  ..., -0.0078, -0.1018, -0.1739],
         [-0.1197, -0.1306, -0.1403,  ..., -0.1600, -0.1617, -0.1827]],

        [[ 0.0168, -0.0307, -0.0600,  ...,  0.0956,  0.0794,  0.0750],
         [-0.1676, -0.1456, -0.1850,  ..., -0.2501, -0.2628, -0.2468],
         [-0.2551, -0.2276, -0.1515,  ..., -0.0386, -0.1205, -0.1908],
         ...,
         [ 0.0117,  0.0067,  0.0812,  ...,  0.1185,  0.1009,  0.1022],
         [-0.0080,  0.0480,  0.0536,  ..., -0.0798, -0.0375, -0.0067],
         [-0.0579, -0.0628, -0.1125,  ..., -0.1455, -0.1393, -0.1766]]],
       grad_fn=<CopySlices>)

***From second batch at conv_tdnn (at return of TDNNbase forward function) Output after optimisation step/***
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, nan, nan],
         [nan, nan, nan,  ..., nan, nan, nan]],

        ...,

        [[nan, nan, nan,  ..., nan, nan, nan],
         [nan, nan, nan,  ..., nan, nan, nan],
         [nan, nan, nan,  ..., nan, nan, nan],
         ...,
         [nan, nan, nan,  ..., nan, nan, nan],
         [nan, nan, nan,  ..., nan, nan, nan],
         [nan, nan, nan,  ..., nan, nan, nan]],

        [[nan, nan, nan,  ..., nan, nan, nan],
         [nan, nan, nan,  ..., nan, nan, nan],
         [nan, nan, nan,  ..., nan, nan, nan],
         ...,
         [nan, nan, nan,  ..., nan, nan, nan],
         [nan, nan, nan,  ..., nan, nan, nan],
         [nan, nan, nan,  ..., nan, nan, nan]],

        [[nan, nan, nan,  ..., nan, nan, nan],
         [nan, nan, nan,  ..., nan, nan, nan],
         [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=<CopySlices>)