chore: set the image used in ptl amp test through set_tf2_image
@@ -9,6 +9,7 @@ def test_pl_mnist() -> None:
exp_dir = "mnist_pl"
config = conf.load_config(conf.cv_examples_path(exp_dir + "/const.yaml"))
config = conf.set_max_length(config, {"batches": 200})
+ config = conf.set_tf2_image(config)
exp.run_basic_test_with_temp_config(config, conf.cv_examples_path(exp_dir), 1)
@@ -18,6 +19,7 @@ def test_pl_mnist_gan() -> None:
exp_dir = "gan_mnist_pl"
config = conf.load_config(conf.gan_examples_path(exp_dir + "/const.yaml"))
config = conf.set_max_length(config, {"batches": 200})
+ config = conf.set_tf2_image(config)
exp.run_basic_test_with_temp_config(config, conf.gan_examples_path(exp_dir), 1)
@@ -28,5 +30,6 @@ def test_pl_const_with_amp(api_style: str) -> None:
dir_name = "pytorch_lightning_amp"
config = conf.load_config(conf.fixtures_path(dir_name + "/" + api_style + "_amp.yaml"))
config = conf.set_max_length(config, {"batches": 200})
+ config = conf.set_tf2_image(config)
exp.run_basic_test_with_temp_config(config, conf.fixtures_path(dir_name), 1)
