torch.load accommodates cpu

Machine-Learning-for-Medical-Language · Sep 13, 2024 · 3b3f3aa · 3b3f3aa
1 parent 9a1c990
commit 3b3f3aa
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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