--train-iters now also overrides time_predictor.num_iterations

main.py

@@ -167,6 +167,7 @@ def run_image_experiment(config: dict, args, dataset_name: str):
     if args.train_iters:
         config["nsgf_training"]["num_iterations"] = args.train_iters
         config["nsf_training"]["num_iterations"] = args.train_iters
+        config["time_predictor"]["num_iterations"] = args.train_iters
     
     # Setup
     data_loader = DatasetLoader(config)