reduced atol

lens_migration.py

@@ -379,6 +379,6 @@ if __name__ == "__main__":
                 log_ps_new = logits_new.log_softmax(-1)
                 log_ps_old = logits_old.log_softmax(-1)
                 print("js div", js_divergence(log_ps_new, log_ps_old))
-                assert (th.allclose(log_ps_new, log_ps_old, atol=1e-7))
+                assert (th.allclose(log_ps_new, log_ps_old, atol=1e-6))
     print("Saving new lens to", args.output_dir)
     tuned_lens.to(th.device("cpu")).save(args.output_dir)