From 1958fc7c7eaf7de3ede3dfddfcde865698c86971 Mon Sep 17 00:00:00 2001 From: Brandon Rozek Date: Thu, 28 Feb 2019 14:41:34 -0500 Subject: [PATCH] Corrected device when constructing fitness tensor --- rltorch/network/ESNetwork.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/rltorch/network/ESNetwork.py b/rltorch/network/ESNetwork.py index c5ae4b7..b3f372c 100644 --- a/rltorch/network/ESNetwork.py +++ b/rltorch/network/ESNetwork.py @@ -3,6 +3,7 @@ import torch from .Network import Network from copy import deepcopy +# [TODO] See if you need to move network to device class ESNetwork(Network): """ Network that functions from the paper Evolutionary Strategies (https://arxiv.org/abs/1703.03864) @@ -52,7 +53,7 @@ class ESNetwork(Network): candidate_solutions = self._generate_candidate_solutions(noise_dict) ## Calculate fitness then mean shift, scale - fitness_values = torch.tensor([self.fitness(x, *args) for x in candidate_solutions]) + fitness_values = torch.tensor([self.fitness(x, *args) for x in candidate_solutions], device = self.device) if self.logger is not None: self.logger.append(self.name + "/" + "fitness_value", fitness_values.mean().item()) fitness_values = (fitness_values - fitness_values.mean()) / (fitness_values.std() + np.finfo('float').eps)