CEL

Code/train.py

@@ -153,12 +153,19 @@ def train(model, name, seed,gradient_accumulation_steps,mixup, threshold, lambda
 #     return loss
 
 #create one hot vector but still keeping the probs
+# def cross_entropy(preds, target):
+#     target = torch.eye(preds.shape[1])[target.long()]
+#     log_preds = torch.log(preds)
+#     loss = -torch.mean(target * log_preds)
+#     return loss
+	
 def cross_entropy(preds, target):
-    target = torch.eye(preds.shape[1])[target.long()]
+    device = preds.device
+    target = torch.eye(preds.shape[1], device=device)[target.long()]
     log_preds = torch.log(preds)
     loss = -torch.mean(target * log_preds)
     return loss
-	
+
 def mixup_function(batch_of_matrices, batch_of_labels, l, t):
 	runs = math.floor(batch_of_matrices.size()[0]/2)
 	counter=0