dear @ptrblck
i was trying to use  OUSMLoss from this repo :  kaggle-panda-challenge-public/noisy_loss.py at master · analokmaus/kaggle-panda-challenge-public · GitHub
for a multi class (5 class) image classification problem.
i was trying to use OUSMLoss like this :
class myOUSMLoss(nn.Module):
    '''
    Implementation of 
    Loss with Online Uncertainty Sample Mining:
    https://arxiv.org/pdf/1901.07759.pdf
    # Params
    k: num of samples to drop in a mini batch
    loss: loss function name (see get_loss function above)
    trigger: the epoch it starts to train on OUSM (please call `.update(epoch)` each epoch)
    '''

    def __init__(self, k=1,  trigger=0, ousm=True):
        super(myOUSMLoss, self).__init__()
        self.k = k
  
        
        self.trigger = trigger
        self.ousm = ousm

    def forward(self, logits, targets, indices=None):
        bs = logits.shape[0]
        print(bs)
        print(bs - self.k)
        if self.ousm and bs - self.k > 0:
            #losses = self.loss(logits, targets, reduction='none')
            losses = bi_tempered_logistic_loss(logits, targets, t1=t1, t2=t2, label_smoothing=smoothing)
            if len(losses.shape) == 2:
                losses = losses.mean(1)
            _, idxs = losses.topk(bs-self.k, largest=False)
            losses = losses.index_select(0, idxs)
            return losses.mean()


and then from train one epoch and valid one epoch function just using loss function like this :
myloss = myOUSMLoss()
loss = myloss(data, target)


and it gives me the following error :

16
14
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-18-d97376da71cb> in <module>()
----> 1 get_ipython().run_cell_magic('time', '', '\nacc_max = 0.\nbest_file = f\'/content/drive/MyDrive/kaggledatasets/cassavamodels/{kernel_type}_best_fold{fold}.pth\'\nfor epoch in range(1, n_epochs+1):\n    print(time.ctime(), \'Epoch:\', epoch)\n    scheduler.step(epoch-1)\n\n    train_loss = train_epoch(train_loader, optimizer)\n    val_loss, acc = val_epoch(valid_loader)\n\n    content = time.ctime() + \' \' + f\'Epoch {epoch}, lr: {optimizer.param_groups[0]["lr"]:.7f}, train loss: {np.mean(train_loss):.5f}, val loss: {np.mean(val_loss):.5f}, acc: {(acc):.5f}\'\n    print(content)\n    with open(f\'/content/drive/MyDrive/kaggledatasets/cassavamodels/log_{kernel_type}.txt\', \'a\') as appender:\n        appender.write(content + \'\\n\')\n\n    if acc > acc_max:\n        print(\'acc ({:.6f} --> {:.6f}).  Saving model ...\'.format(acc_max, acc))\n        torch.save(model.state_dict(), best_file)\n        acc_max = acc\n\n\n#torch.save(model.state_dict(), os.path.join(f\'/content/drive/MyDrive/kaggledatasets/cassavamodels/{kernel_type}_final_fold{fold}.pth\'))')
5 frames
<decorator-gen-60> in time(self, line, cell, local_ns)
<timed exec> in <module>()
<ipython-input-14-5055b70e47f6>
 in forward(self, logits, targets, indices)
     31             if len(losses.shape) == 2:
     32                 losses = losses.mean(1)
---> 33             _, idxs = losses.topk(bs-self.k, largest=False)
     34             losses = losses.index_select(0, idxs)
     35             return losses.mean()
RuntimeError: invalid argument 5: k not in range for dimension at /pytorch/aten/src/THC/generic/THCTensorTopK.cu:26

how can i fix this issue?