Merge pull request #66 from traja-team/fortasyn

fixing accuracy > 1 error
traja-team · Feb 13, 2021 · f87978d · f87978d
2 parents 045a4bd + 51f7c0e
commit f87978d
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/traja/models/losses.py b/traja/models/losses.py
@@ -1,5 +1,6 @@
 import torch
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
 class Criterion:
     """Implements the loss functions of Autoencoders, Variational Autoencoders and LSTM models
@@ -51,7 +52,8 @@ def classifier_criterion(self, predicted, target):
         :return: Cross entropy loss
         """
 
-        # _, predicted = torch.max(predicted.data, 1)
+        predicted = predicted.to(device)
+        target = target.to(device)
         loss = self.crossentropy_loss(predicted, target.view(-1))
         return loss
 

diff --git a/traja/models/train.py b/traja/models/train.py
@@ -304,7 +304,8 @@ def fit(
                             # Compute number of correct samples
                             total += ids.size(0)
                             _, predicted = torch.max(classifier_out.data, 1)
-                            correct += (predicted == classes).sum().item()
+
+                            correct += (predicted.cpu() == classes.cpu().T).sum().item()
 
                         if self.regress:
                             regressor_out = self.model(
@@ -401,7 +402,7 @@ def validate(self, validation_loader):
                     # Compute number of correct samples
                     total += ids.size(0)
                     _, predicted = torch.max(classifier_out.data, 1)
-                    correct += (predicted == classes).sum().item()
+                    correct += (predicted.cpu() == classes.cpu().T).sum().item()
 
                 if self.regress:
                     regressor_out = self.model(