From 923ba8978aa06ac487dfd37b16d1862634b47bc0 Mon Sep 17 00:00:00 2001 From: "Dr. Kashif Rasul" Date: Wed, 12 Feb 2020 17:18:02 +0100 Subject: [PATCH] move predictor network to device --- pts/model/simple_feedforward/simple_feedforward_estimator.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pts/model/simple_feedforward/simple_feedforward_estimator.py b/pts/model/simple_feedforward/simple_feedforward_estimator.py index 50adda2..87c6b22 100644 --- a/pts/model/simple_feedforward/simple_feedforward_estimator.py +++ b/pts/model/simple_feedforward/simple_feedforward_estimator.py @@ -158,7 +158,7 @@ class SimpleFeedForwardEstimator(PTSEstimator): batch_normalization=self.batch_normalization, mean_scaling=self.mean_scaling, num_parallel_samples=self.num_parallel_samples, - ) + ).to(device) copy_parameters(trained_network, prediction_network)