diff --git a/src/asleep/models.py b/src/asleep/models.py
index f7dfc76..6e5360d 100644
--- a/src/asleep/models.py
+++ b/src/asleep/models.py
@@ -193,7 +193,7 @@ def predict(self, X, groups=None):
         model.to(self.device)
 
         _, y_pred, _ = sslmodel.predict(
-            model, dataloader, self.device, output_logits=False)
+            model, dataloader, self.device, output_logits=False, name='prediction')
 
         y_pred = self.hmms.predict(y_pred, groups=groups)
 
diff --git a/src/asleep/sslmodel.py b/src/asleep/sslmodel.py
index 4bae9d0..5684e52 100644
--- a/src/asleep/sslmodel.py
+++ b/src/asleep/sslmodel.py
@@ -115,7 +115,8 @@ def get_sslnet(tag='v1.0.0', pretrained=False):
     return sslnet
 
 
-def predict(model, data_loader, device, output_logits=False):
+def predict(model, data_loader, device,
+            output_logits=False, name='train'):
     """
     Iterate over the dataloader and do prediction with a pytorch model.
     :param nn.Module model: pytorch Module
@@ -145,8 +146,12 @@ def predict(model, data_loader, device, output_logits=False):
                 pred_y = torch.argmax(logits, dim=1)
                 predictions_list.append(pred_y.cpu())
             pid_list.extend(pid)
-    true_list = torch.cat(true_list)
+
     predictions_list = torch.cat(predictions_list)
+    if name == 'prediction':
+        true_list = predictions_list
+    else:
+        true_list = torch.Tensor([1, 2, 3])
 
     if output_logits:
         return (