From 3b3f3aa5b7d79824f42804f447f795e9ac605ba7 Mon Sep 17 00:00:00 2001 From: Spencer Thomas Date: Fri, 13 Sep 2024 13:06:41 -0400 Subject: [PATCH] torch.load accommodates cpu --- src/cnlpt/api/cnn_rest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cnlpt/api/cnn_rest.py b/src/cnlpt/api/cnn_rest.py index 93d3bdd..e0fffb2 100644 --- a/src/cnlpt/api/cnn_rest.py +++ b/src/cnlpt/api/cnn_rest.py @@ -61,7 +61,7 @@ async def startup_event(): embed_dims = conf_dict['cnn_embed_dim'], num_filters=conf_dict['num_filters'], filters=conf_dict['filters']) - model.load_state_dict(torch.load(join(model_name, 'pytorch_model.bin')), map_location=torch.device(device)) + model.load_state_dict(torch.load(join(model_name, 'pytorch_model.bin'), map_location=torch.device(device))) app.state.model = model.to(device) app.state.tokenizer = tokenizer