Torch approximator now uses correct device

MushroomRL · Dec 4, 2023 · 42872ec · 42872ec
1 parent 74c4bf5
commit 42872ec
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/mushroom_rl/approximators/parametric/torch_approximator.py b/mushroom_rl/approximators/parametric/torch_approximator.py
@@ -49,6 +49,7 @@ def __init__(self, input_shape, output_shape, network, optimizer=None, loss=None
         self._n_fit_targets = n_fit_targets
 
         self.network = network(input_shape, output_shape, dropout=dropout, **params)
+        self.network.to(TorchUtils.get_device())
 
         if self._dropout:
             self.network.eval()