Skip to content

Commit

Permalink
Use absl for neuralgcm
Browse files Browse the repository at this point in the history
  • Loading branch information
wsmoses committed Nov 8, 2024
1 parent b25b8c0 commit 5dd9576
Showing 1 changed file with 5 additions and 2 deletions.
7 changes: 5 additions & 2 deletions test/neuralgcm_test.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
from absl import app
import jax.numpy as jnp
import jax.random
import jax.lax
Expand Down Expand Up @@ -178,8 +179,10 @@ def run_on_fn(self, fn, steps=1):
},
).timeit(steps)


if __name__ == "__main__":
def main():
c = NeuralGCM()
c.setUp()
c.test()

if __name__ == "__main__":
app.run(main)

0 comments on commit 5dd9576

Please sign in to comment.