wassname/pytorch-ts
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-ts.git synced 2026-07-12 20:25:39 +08:00
Code Issues Packages Projects Releases Wiki Activity

add cross platform training and inference (#27)

Browse Source 
Adapt model deserialization to support between cpu training and gpu inference and viceversa
This commit is contained in:
Edrin Basha
2020-12-29 17:06:13 +01:00
committed by GitHub Enterprise
parent ecc31f6082
commit 1be0a7b305
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pts/model/predictor.py
+1 -1
View File
@@ -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: