@ptrblck  fixed that issue and now getting Nan
this is the code i am using :
def soft_cross_entropy_loss(logits, targets, weights=1, reduction='none'):

    if len(targets.shape) == 1 or targets.shape[1] == 1:

        onehot_targets = torch.eye(logits.shape[1])[targets].to(logits.device)

    else:

        onehot_targets = targets

    loss = -torch.sum(onehot_targets * F.log_softmax(logits, 1), 1)

    if reduction == 'none':

        return loss

    elif reduction == 'sum':

        return loss.sum()

    elif reduction == 'mean':

        return loss.mean()

def ousm(logits, targets, indices=None):

    logits, targets = _check_input_type(logits, targets, 'SoftCrossEntropy')

    bs = logits.shape[0]

    k = 1

    print(bs)

    print("bs - k ", bs - k)

    

    if bs - k > 0:

      #losses = self.loss(logits, targets)

      losses = soft_cross_entropy_loss(logits, targets, reduction='none')

      #losses = bi_tempered_logistic_loss(logits, targets, t1=t1, t2=t2, label_smoothing=smoothing)

      print(losses)

      print("losses.shape ",losses.shape)

      print("losses ",list(losses.size()))

      print(len(losses.shape))

      if len(losses.shape) == 2:

          losses = losses.mean(1)

      print("losses.shape test ",losses.shape)

      print("losses test",list(losses.size()))

      _, idxs = losses.topk(bs-k, largest=True)

      losses = losses.index_select(0, idxs)

      return losses.mean()

and i get output like this :

16
bs - k  15
tensor([1.5355, 1.6762, 1.6320, 1.5614, 1.5086, 1.6454, 1.4942, 1.5916, 1.6462,
        1.7642, 1.4607, 1.5639, 1.6708, 1.4343, 1.7360, 1.6173],
       device='cuda:0', grad_fn=<NegBackward>)
losses.shape  torch.Size([16])
losses  [16]
1
losses.shape test  torch.Size([16])
losses test [16]

after some times,
16
bs - k  15
tensor([nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan, nan],
       device='cuda:0', grad_fn=<NegBackward>)
losses.shape  torch.Size([16])
losses  [16]
1
losses.shape test  torch.Size([16])
losses test [16]