Made sure everything went to their appropriate devices

rltorch/action_selector/ArgMaxSelector.py

@@ -10,7 +10,7 @@ class ArgMaxSelector:
     def best_act(self, state):
         with torch.no_grad():
             if self.device is not None:
-                self.device.to(self.device)
+                state = state.to(self.device)
             action_values = self.model(state).squeeze(0)
             action = self.random_act() if (action_values[0] == action_values).all() else action_values.argmax().item()
         return action

rltorch/network/Network.py

@@ -2,11 +2,14 @@ class Network:
     """
     Wrapper around model which provides copy of it instead of trained weights
     """
-    def __init__(self, model, optimizer, config, logger = None, name = ""):
+    def __init__(self, model, optimizer, config, device = None, logger = None, name = ""):
         self.model = model
         self.optimizer = optimizer(model.parameters(), lr = config['learning_rate'], weight_decay = config['weight_decay'])
         self.logger = logger
         self.name = name
+        self.device = device
+        if self.device is not None:
+            self.model = self.model.to(device)
 
     def __call__(self, *args):
         return self.model(*args)

rltorch/network/TargetNetwork.py

@@ -4,9 +4,11 @@ class TargetNetwork:
     """
     Wrapper around model which provides copy of it instead of trained weights
     """
-    def __init__(self, network):
+    def __init__(self, network, device = None):
         self.model = network.model
         self.target_model = deepcopy(network.model)
+        if network.device is not None:
+            self.target_model = self.target_model.to(network.device)
 
     def __call__(self, *args):
         return self.model(*args)