diff --git a/test/neuralgcm_test.py b/test/neuralgcm_test.py
index 7642f60a..e7f822ed 100644
--- a/test/neuralgcm_test.py
+++ b/test/neuralgcm_test.py
@@ -180,7 +180,7 @@ def run_on_fn(self, fn, steps=1):
         ).timeit(steps)
 
 
-def main():
+def main(argv):
     c = NeuralGCM()
     c.setUp()
     c.test()