Repair OCR RuntimeError: Expected all tensors to be on the same devic…

…e, but found at least two devices, cuda:0 and cpu!
ria-com · Jan 6, 2025 · f1ee5d4 · f1ee5d4
1 parent b279baa
commit f1ee5d4
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/nomeroff_net/nnmodels/ocr_model.py b/nomeroff_net/nnmodels/ocr_model.py
@@ -156,10 +156,10 @@ def calculate_loss(self, logits, texts):
         logits_lens = torch.full(size=(batch_size,), fill_value=input_len, dtype=torch.int32)
         # calculate ctc
         loss = self.criterion(
-            logits,
-            encoded_texts,
+            logits.to(device),
+            encoded_texts.to(device),
             logits_lens.to(device),
-            text_lens)
+            text_lens.to(device))
         return loss
 
     def step(self, batch):

diff --git a/nomeroff_net/pipes/number_plate_text_readers/base/ocr.py b/nomeroff_net/pipes/number_plate_text_readers/base/ocr.py
@@ -382,10 +382,10 @@ def get_acc(self, predicted: List, decode: List) -> torch.Tensor:
         logits_lens = torch.full(size=(batch_size,), fill_value=input_len, dtype=torch.int32)
 
         acc = functional.ctc_loss(
-            logits,
-            encoded_texts,
+            logits.to(device),
+            encoded_texts.to(device),
             logits_lens.to(device),
-            text_lens
+            text_lens.to(device)
         )
         return 1 - acc / len(self.letters)