reverse

requirements.txt

@@ -14,7 +14,6 @@ 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
 

src/evaluation.py

@@ -15,9 +15,6 @@ 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",
@@ -353,7 +350,7 @@ def evaluate(model_id, revision):
         observations, _ = envs.reset()
         episodic_returns = []
         while len(episodic_returns) < 10:
-            actions = agent(torch.tensor(observations, device="cuda")).cpu().numpy()
+            actions = agent(torch.tensor(observations)).numpy()
             observations, _, _, _, infos = envs.step(actions)
             if "final_info" in infos:
                 for info in infos["final_info"]: