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

use device

Browse Source
This commit is contained in:
Kashif Rasul
2019-12-01 10:07:00 +01:00
parent 7905888848
commit 78ebf59d15
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pts/model/deepar/deepar_estimator.py
+2 -2
View File
@@ -151,7 +151,7 @@ class DeepAREstimator(PTSEstimator):
),
])
def create_training_network(self) -> DeepARTrainingNetwork:
def create_training_network(self, device: torch.device) -> DeepARTrainingNetwork:
return DeepARTrainingNetwork(
num_layers=self.num_layers,
num_cells=self.num_cells,
@@ -166,4 +166,4 @@ class DeepAREstimator(PTSEstimator):
lags_seq=self.lags_seq,
scaling=self.scaling,
dtype=self.dtype,
)
).to(device)