From 78ebf59d15fbb5400b0bc471438cae4154f7ead4 Mon Sep 17 00:00:00 2001
From: Kashif Rasul <kashif.rasul@gmail.com>
Date: Sun, 1 Dec 2019 10:07:00 +0100
Subject: [PATCH] use device

---
 pts/model/deepar/deepar_estimator.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pts/model/deepar/deepar_estimator.py b/pts/model/deepar/deepar_estimator.py
index 6b741a1..cf3b33d 100644
--- a/pts/model/deepar/deepar_estimator.py
+++ b/pts/model/deepar/deepar_estimator.py
@@ -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)