dimension

Code/train.py

@@ -135,7 +135,7 @@ def train(model, name, seed,gradient_accumulation_steps,mixup, threshold, lambda
 	return evaluation_test, evaluation_train
 
 def SoftCrossEntropyLoss (input, target):
-    logprobs = torch.nn.functional.log_softmax (input, dim = 2)
+    logprobs = torch.nn.functional.log_softmax (input, dim = 0)
     return  -(target * logprobs).sum() / input.shape[0]
 	
 def mixup_function(batch_of_matrices, batch_of_labels, l, t):