Expose device arg

Author: MartinBraquet

squadro/agents/montecarlo_agent.py

@@ -470,6 +470,7 @@ def __init__(
         self,
         model_path: str = None,
         model_config=None,
+        device: str = None,
         mcts_kwargs: Optional[dict] = None,
         **kwargs,
     ):
@@ -478,7 +479,8 @@ def __init__(
         if 'evaluator' not in kwargs:
             kwargs['evaluator'] = DeepQLearningEvaluator(
                 model_path=model_path,
-                model_config=model_config
+                model_config=model_config,
+                device=device,
             )
         super().__init__(mcts_kwargs=mcts_kwargs, **kwargs)
 

squadro/training/deep_q_learning.py

@@ -815,6 +815,7 @@ def __init__(
         adaptive_sampling=True,
         freeze_backprop=None,
         display_plot=True,
+        device=None,
     ):
         """
         Represents a class for initializing and training a deep reinforcement learning agent
@@ -854,6 +855,7 @@ def __init__(
         self.agent_kwargs = dict(
             mcts_kwargs=self.mcts_kwargs,
             max_time_per_move=inf,
+            device=device,
         )
 
         self.agent = MonteCarloDeepQLearningAgent(