From a75ff3fd670abba66e43475096a1c4eae864f9ad Mon Sep 17 00:00:00 2001 From: Navin Soni Date: Mon, 27 Jun 2022 21:58:30 +0000 Subject: [PATCH] change: Update model name from 'compiled.pt' to 'model.pth' for neo --- tests/data/pytorch_neo/code/inference.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/data/pytorch_neo/code/inference.py b/tests/data/pytorch_neo/code/inference.py index 984005039f..5b89c2bebc 100644 --- a/tests/data/pytorch_neo/code/inference.py +++ b/tests/data/pytorch_neo/code/inference.py @@ -71,8 +71,8 @@ def model_fn(model_dir): logger.info("model_fn") neopytorch.config(model_dir=model_dir, neo_runtime=True) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - # The compiled model is saved as "compiled.pt" - model = torch.jit.load(os.path.join(model_dir, "compiled.pt"), map_location=device) + # The compiled model is saved as "model.pth" + model = torch.jit.load(os.path.join(model_dir, "model.pth"), map_location=device) # It is recommended to run warm-up inference during model load sample_input_path = os.path.join(model_dir, "sample_input.pkl")