From 438785cde164805877c1eae0295c650f88fe4013 Mon Sep 17 00:00:00 2001 From: "Dr. Kashif Rasul" Date: Tue, 28 Jan 2020 12:12:32 +0100 Subject: [PATCH] copy trainer network to device --- pts/model/transformer/transformer_estimator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pts/model/transformer/transformer_estimator.py b/pts/model/transformer/transformer_estimator.py index 6d3a6a1..2a86d36 100644 --- a/pts/model/transformer/transformer_estimator.py +++ b/pts/model/transformer/transformer_estimator.py @@ -215,7 +215,7 @@ class TransformerEstimator(PTSEstimator): lags_seq=self.lags_seq, scaling=self.scaling, num_parallel_samples=self.num_parallel_samples, - ) + ).to(device) copy_parameters(trained_network, prediction_network)