Fetching metadata from the HF Docker repository...
|
@@ -1,5 +1,5 @@
|
|
1
|
import os
|
|
2
|
-
import jax
|
|
3
|
import jax.numpy as jnp
|
|
4
|
import flax.linen as nn
|
|
5
|
import pickle
|
|
|
|
1
|
import os
|
|
2
|
+
import Jax[CPU]
|
|
3
|
import jax.numpy as jnp
|
|
4
|
import flax.linen as nn
|
|
5
|
import pickle
|
