update test

Author: JackCaoG

test/run_tests.sh

@@ -93,7 +93,9 @@ function run_all_tests {
   run_opbyop python3 "$CDIR/test_operations.py" "$@" --verbosity=$VERBOSITY
   run_eager_debug python3 "$CDIR/test_operations.py" "$@" --verbosity=$VERBOSITY
   run_async_rng python3 "$CDIR/test_operations.py" "$@" --verbosity=$VERBOSITY
-  run_test python3 "$CDIR/test_checkpoint.py"
+  # TODO: enable this test after tf update, currently optimization_barrier does not
+  # work on CPU.
+  # run_test python3 "$CDIR/test_checkpoint.py"
   run_test python3 "$CDIR/test_mp_replication.py"
   run_test python3 "$CDIR/test_mp_all_to_all.py"
   run_test python3 "$CDIR/test_mp_collective_permute.py"

test/test_checkpoint.py

@@ -5,7 +5,7 @@
 import torch_xla.utils.checkpoint as checkpoint
 
 
-def run(grad_checkpoint):
+def run():
   device = xm.xla_device()
   model = torch.nn.ModuleList([
       torch.nn.Sequential(
@@ -22,16 +22,15 @@ def run(grad_checkpoint):
     optimizer.zero_grad()
     x = dummy_data
     for n_l, layer in enumerate(model):
-      x = checkpoint.checkpoint(layer, x)
+      if n_l > 0:
+        x = checkpoint.checkpoint(layer, x)
+      else:
+        x = layer(x)
     dummy_loss = x.sum()
     dummy_loss.backward()
     optimizer.step()
     xm.mark_step()
-    xm.wait_device_ops()
 
 
 if __name__ == "__main__":
-  parser = argparse.ArgumentParser()
-  parser.add_argument("--grad_checkpoint", type=int, required=True)
-  args = parser.parse_args()
-  run(args.grad_checkpoint)
+  run()