From c4e8ab5e43dfaed45ee2d9044c523c371c3d9e4d Mon Sep 17 00:00:00 2001 From: bkim Date: Wed, 14 Aug 2024 14:15:36 +0900 Subject: [PATCH] fix FutureWarning in torch.load, since torch 2.3.1 --- easyocr/detection.py | 2 +- easyocr/recognition.py | 2 +- unit_test/README.md | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/easyocr/detection.py b/easyocr/detection.py index 072178a23..5a04c1716 100644 --- a/easyocr/detection.py +++ b/easyocr/detection.py @@ -82,7 +82,7 @@ def get_detector(trained_model, device='cpu', quantize=True, cudnn_benchmark=Fal except: pass else: - net.load_state_dict(copyStateDict(torch.load(trained_model, map_location=device))) + net.load_state_dict(copyStateDict(torch.load(trained_model, map_location=device, weights_only=True))) net = torch.nn.DataParallel(net).to(device) cudnn.benchmark = cudnn_benchmark diff --git a/easyocr/recognition.py b/easyocr/recognition.py index 147370f7d..bd8dd9b72 100644 --- a/easyocr/recognition.py +++ b/easyocr/recognition.py @@ -179,7 +179,7 @@ def get_recognizer(recog_network, network_params, character,\ pass else: model = torch.nn.DataParallel(model).to(device) - model.load_state_dict(torch.load(model_path, map_location=device)) + model.load_state_dict(torch.load(model_path, map_location=device, weights_only=True)) return model, converter diff --git a/unit_test/README.md b/unit_test/README.md index f88bd75ff..d37d1a166 100644 --- a/unit_test/README.md +++ b/unit_test/README.md @@ -9,7 +9,7 @@ This module can be used as a typical python module. One python wrapper script an ### Python script (*recommneded*) The script can be called with (assuming calling from `EasyOCR/`); ``` -python ./unit_test/run_unit_test.py --easyocr ./easyocr --verbose 2 --test ./unit_test/EasyOcrUnitTestPackage.pickle --data_dir ./examples +python ./unit_test/run_unit_test.py --easyocr ./easyocr --verbose 2 --test_data ./unit_test/data/EasyOcrUnitTestPackage.pickle --image_data_dir ./examples ``` #### Script arguments