.div

Code/train.py

@@ -146,7 +146,8 @@ def train(model, name, seed,gradient_accumulation_steps,mixup, threshold, lambda
 #pytorch Forum, try with dim=1 for targets
 def cross_entropy(logits, target):
     logprobs = torch.nn.functional.log_softmax (logits, dim = 1)
-    return  -torch.sum(target * logprobs, dim=1) / logits.shape[0]
+    loss = -torch.sum(target * logprobs, dim=1)
+    return loss.div(logits.shape[0])
 
 #gives average #2
 # def cross_entropy(preds, target):