From 1be0a7b305ba78cd8f2353ef1ff70fd3cc7dc286 Mon Sep 17 00:00:00 2001 From: Edrin Basha Date: Tue, 29 Dec 2020 17:06:13 +0100 Subject: [PATCH] add cross platform training and inference (#27) Adapt model deserialization to support between cpu training and gpu inference and viceversa --- pts/model/predictor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pts/model/predictor.py b/pts/model/predictor.py index a232705..75f8c56 100644 --- a/pts/model/predictor.py +++ b/pts/model/predictor.py @@ -175,7 +175,7 @@ class PTSPredictor(Predictor): model_name = 'prediction_net' with (path / f"{model_name}-network.json").open("r") as fp: prediction_net = load_json(fp.read()) - prediction_net.load_state_dict(torch.load(path / "prediction_net")) + prediction_net.load_state_dict(torch.load(path / "prediction_net", map_location=device)) # input_names is derived from the prediction_net if "input_names" in parameters: