diff --git a/scripts/train.py b/scripts/train.py
index 944956b..492c887 100644
--- a/scripts/train.py
+++ b/scripts/train.py
@@ -2,8 +2,6 @@ import subprocess
 import time
 
 import torch
-import tyro
-import yaml
 import os
 
 from brittle_star_project.dataclasses import PPOArgs
diff --git a/src/brittle_star_project/trainers/PPOTrainer.py b/src/brittle_star_project/trainers/PPOTrainer.py
index d4297de..8c238f1 100644
--- a/src/brittle_star_project/trainers/PPOTrainer.py
+++ b/src/brittle_star_project/trainers/PPOTrainer.py
@@ -5,7 +5,6 @@ from dataclasses import asdict, dataclass
 from functools import partial
 from typing import Any
 
-import flax
 import jax
 import jax.numpy as jnp
 import numpy as np
diff --git a/src/experiment_logger/wandb_utils.py b/src/experiment_logger/wandb_utils.py
index 2c162fb..302d308 100644
--- a/src/experiment_logger/wandb_utils.py
+++ b/src/experiment_logger/wandb_utils.py
@@ -52,7 +52,8 @@ def init_wandb(
 
         if not has_key and not is_interactive and os.environ.get("WANDB_MODE") != "offline":
             logger.warning(
-                "WANDB_API_KEY not found and environment is non-interactive. Switching to offline mode."
+                "WANDB_API_KEY not found and environment is non-interactive. "
+                "Switching to offline mode."
             )
             sync_path = f"runs/{name}" if name else "runs"
             logger.warning(f"WandB is offline. Use 'wandb sync {sync_path}' to upload logs later.")