try cuda

requirements.txt

@@ -14,9 +14,11 @@ pandas==2.0.0
 python-dateutil==2.8.2
 requests==2.28.2
 rliable==1.0.8
+--extra-index-url https://download.pytorch.org/whl/cu113
 torch==2.2.2
 tqdm==4.65.0
 
+
 # Log Visualizer
 BeautifulSoup4==4.12.2
 lxml==4.9.3

src/evaluation.py

@@ -15,6 +15,10 @@ logger = setup_logger(__name__)
 
 API = HfApi(token=os.environ.get("TOKEN"))
 
+logger.info(f"Is CUDA available: {torch.cuda.is_available()}")
+logger.info(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
+
+
 ALL_ENV_IDS = [
     "AdventureNoFrameskip-v4",
     "AirRaidNoFrameskip-v4",
@@ -338,7 +342,7 @@ def evaluate(model_id, revision):
 
     # Load the agent
     try:
-        agent = torch.jit.load(agent_path)
+        agent = torch.jit.load(agent_path).to("cuda")
     except Exception as e:
         logger.error(f"Error loading agent: {e}")
         return None
@@ -349,7 +353,7 @@ def evaluate(model_id, revision):
         observations, _ = envs.reset()
         episodic_returns = []
         while len(episodic_returns) < 10:
-            actions = agent(torch.tensor(observations)).numpy()
+            actions = agent(torch.tensor(observations, device="cuda")).cpu().numpy()
             observations, _, _, _, infos = envs.step(actions)
             if "final_info" in infos:
                 for info in infos["final_info"]: