Hey @albanD,
Apologies for the direct message.
If I am using a dynamically created tensor for loss calculation, what would be the recommended approach for writing the loss function for optimizing Cuda command use? Ex. -
class KLDLoss(nn.Module):
    """KL-divergence loss between attention weight and uniform distribution"""
    def __init__(self):
        super(KLDLoss, self).__init__()
        
    def forward(self, attn_val, cluster):
        """
        Example:
          Input - attention value = torch.tensor([0.05, 0.1, 0.05, 0.1, 0.05, 0.1, 0.05, 0.05, 
                                0.1, 0.05, 0.1, 0.05, 0.1, 0.05])
                  cluster = [1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2]
          Output - 0.0043
        """
        
        kld_loss = 0
        cluster = np.array(cluster)
        for cls in np.unique(cluster):
            index = np.where(cluster==cls)[0]
            if torch.cuda.is_available():
                kld_loss += F.kl_div(F.log_softmax(attn_val[index], dim=0)[None],\
                                torch.ones(len(index), 1)[None].cuda()/len(index),\
                                reduction = 'batchmean')
            else:
                kld_loss += F.kl_div(F.log_softmax(attn_val[index], dim=0)[None],\
                                torch.ones(len(index), 1)[None]/len(index),\
                                reduction = 'batchmean')                
        return kld_loss

Thanks!