Gradients being tiny, I try to use nn.utils.clip_grad_value_(MODEL.parameters(), 1) but it didnt. But i am thinking, If model architecture was my problem, but I’ve been converting from another framework, which look similar,
def Siamese(vocab_size=len(vocab), d_model=128, mode='train'):
    def normalize(x):  # normalizes the vectors to have L2 norm normalize
        return x / fastnp.sqrt(fastnp.sum(x * x, axis=-1, keepdims=True))
    
    q_processor = tl.Serial(  # Processor will run on Q1 and Q2.
        tl.Embedding(vocab_size, d_model), # Embedding layer
        tl.LSTM(d_model), # LSTM layer
        tl.Mean(axis=1), # Mean over columns
        tl.Fn('Normalize', lambda x: normalize(x))  # Apply normalize function
    ) 
    # Run on Q1 and Q2 in parallel.
    model = tl.Parallel(q_processor, q_processor)
    return model

which is same as our pytorch architecture
    def forward(self, x1, x2):

        lstm_out_1, _ = self.em(x1)
        mean_layer_1 = lstm_out_1.mean(1)
        normalized_1 = F.normalize(mean_layer_1)

        lstm_out_2, _ = self.em(x2)
        mean_layer_2 = lstm_out_2.mean(1)
        normalized_2 = F.normalize(mean_layer_2)
        
        return normalized_1, normalized_2

In case of loss , I also tried to use other distance losses such as torch.norm or torch.dist to see if grad is giving good, but gave same tiny result , It look same as before.  Here is that framework hardtriplet loss but My pytorch loss is right as this, and gave same result when testing .


Summary
def TripletLoss(v1, v2, margin=0.25):
    scores = fastnp.dot(v1, v2.T)  # pairwise cosine sim
    batch_size = len(scores)
    positive = fastnp.diagonal(scores) # the positive ones (duplicates)
    negative_without_positive = scores - 2.0 * fastnp.eye(batch_size)
    closest_negative = negative_without_positive.max(axis=1)
    negative_zero_on_duplicate =  scores * (1.0 - fastnp.eye(batch_size))
    mean_negative = fastnp.sum(negative_zero_on_duplicate, axis=1) / (batch_size - 1)
    triplet_loss1 = fastnp.maximum(0.0, margin - positive + closest_negative)
    triplet_loss2 = fastnp.maximum(0.0, margin - positive + mean_negative)
    triplet_loss = fastnp.mean(triplet_loss1 + triplet_loss2)
    return triplet_loss


Training this trax framework was working and gave good results but converting to pytorch I’m stuck  but cant figure out, I am thinking where could be my problem because using different architecture also didnt solve problem of grad to be tiny or zero.Help!!
That None was not expected in my gradient, using different architecture  removed but came all zeros instead.