From 6183d1f8fc07508e5ebc8e2c2bee2298cba1651c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pedro=20Ribeiro=20Mendes=20J=C3=BAnior?= Date: Sun, 27 Apr 2025 19:24:20 -0300 Subject: [PATCH] do not pin memory when device is CPU to avoid warning --- easyocr/recognition.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/easyocr/recognition.py b/easyocr/recognition.py index 530ef9517e2..a3b74789c78 100644 --- a/easyocr/recognition.py +++ b/easyocr/recognition.py @@ -200,7 +200,7 @@ def get_text(character, imgH, imgW, recognizer, converter, image_list,\ test_data = ListDataset(img_list) test_loader = torch.utils.data.DataLoader( test_data, batch_size=batch_size, shuffle=False, - num_workers=int(workers), collate_fn=AlignCollate_normal, pin_memory=True) + num_workers=int(workers), collate_fn=AlignCollate_normal, pin_memory=device != 'cpu') # predict first round result1 = recognizer_predict(recognizer, converter, test_loader,batch_max_length,\ @@ -214,7 +214,7 @@ def get_text(character, imgH, imgW, recognizer, converter, image_list,\ test_data = ListDataset(img_list2) test_loader = torch.utils.data.DataLoader( test_data, batch_size=batch_size, shuffle=False, - num_workers=int(workers), collate_fn=AlignCollate_contrast, pin_memory=True) + num_workers=int(workers), collate_fn=AlignCollate_contrast, pin_memory=device != 'cpu') result2 = recognizer_predict(recognizer, converter, test_loader, batch_max_length,\ ignore_idx, char_group_idx, decoder, beamWidth, device = device)