feat(train): local jax cache

borisdayma · borisdayma · commit 9f5e879417e7 · 2022-02-23T00:56:05.000Z
diff --git a/tools/train/train.py b/tools/train/train.py
@@ -57,9 +57,7 @@
     set_partitions,
 )
 
-cc.initialize_cache(
-    "/home/boris/dalle-mini/jax_cache", max_cache_size_bytes=5 * 2**30
-)
+cc.initialize_cache("./jax_cache", max_cache_size_bytes=5 * 2**30)
 
 logger = logging.getLogger(__name__)
 

Original file line number	Diff line number	Diff line change
`@@ -57,9 +57,7 @@`
`57`	`57`	`set_partitions,`
`58`	`58`	`)`
`59`	`59`
`60`		`-cc.initialize_cache(`
`61`		`- "/home/boris/dalle-mini/jax_cache", max_cache_size_bytes=5 * 2**30`
`62`		`-)`
	`60`	`+cc.initialize_cache("./jax_cache", max_cache_size_bytes=5 * 2**30)`
`63`	`61`
`64`	`62`	`logger = logging.getLogger(__name__)`
`65`	`63`