From 5c7b35d6948ed9f9840c3e45ae9cdef4562e02de Mon Sep 17 00:00:00 2001 From: Sven Mika Date: Thu, 17 Sep 2020 09:04:46 +0200 Subject: [PATCH] [RLlib] Issue 10833 TorchPolicy GPU. (#10834) --- rllib/policy/torch_policy.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/rllib/policy/torch_policy.py b/rllib/policy/torch_policy.py index f5dd7be39..2521afda9 100644 --- a/rllib/policy/torch_policy.py +++ b/rllib/policy/torch_policy.py @@ -4,7 +4,6 @@ import numpy as np import time from typing import Callable, Dict, List, Optional, Tuple, Type, Union -import ray from ray.rllib.models.modelv2 import ModelV2 from ray.rllib.models.torch.torch_modelv2 import TorchModelV2 from ray.rllib.models.torch.torch_action_dist import TorchDistributionWrapper @@ -104,7 +103,7 @@ class TorchPolicy(Policy): """ self.framework = "torch" super().__init__(observation_space, action_space, config) - if torch.cuda.is_available() and ray.get_gpu_ids(): + if torch.cuda.is_available(): self.device = torch.device("cuda") else: self.device = torch.device("cpu")