tests/test_prediction.py

 # -*- coding: utf-8 -*-
 
 import pytest
-import torch
 
 from dan.predict.prediction import DAN
 from dan.utils import read_image
@@ -46,9 +45,9 @@ def test_predict(
     image = read_image(prediction_data_path / "images" / image_name)
     image = dan_model.preprocess(image)
 
-    input_tensor = torch.tensor(image).permute(2, 0, 1).unsqueeze(0)
+    input_tensor = image.unsqueeze(0)
     input_tensor = input_tensor.to(device)
-    input_sizes = [image.shape[:2]]
+    input_sizes = [image.shape[1:]]
 
     prediction = dan_model.predict(input_tensor, input_sizes)
 

tests/test_training.py

@@ -25,7 +25,7 @@ from tests.conftest import FIXTURES
             },
             {
                 "nb_chars": 41,
-                "cer": 1.2683,
+                "cer": 1.2927,
                 "nb_words": 9,
                 "wer": 1.0,
                 "nb_words_no_punct": 9,
@@ -34,7 +34,7 @@ from tests.conftest import FIXTURES
             },
             {
                 "nb_chars": 49,
-                "cer": 1.1224,
+                "cer": 1.102,
                 "nb_words": 9,
                 "wer": 1.0,
                 "nb_words_no_punct": 9,