make more min

kscalelabs · Aug 20, 2024 · 215bff0 · 215bff0
1 parent 36b8aae
commit 215bff0
Show file tree

Hide file tree

Showing 3 changed files with 1 addition and 133 deletions.
diff --git a/mnist.py b/mnist.py
diff --git a/train.py b/train.py
@@ -4,7 +4,6 @@
 import logging
 import os
 from dataclasses import dataclass, field
-from functools import partial
 from typing import Any, Dict, List, Tuple
 
 import equinox as eqx
@@ -13,12 +12,12 @@
 import mediapy as media
 import numpy as np
 import optax
+import wandb
 from brax.envs import State
 from brax.mjx.base import State as MjxState
 from jax import Array
 from tqdm import tqdm
 
-import wandb
 from environment import HumanoidEnv
 
 logger = logging.getLogger(__name__)

diff --git a/training.notes b/training.notes