Update train.py

Revert changes to the corrdiff train examples.

examples/generative/corrdiff/train.py

@@ -268,14 +268,8 @@ def main(cfg: DictConfig) -> None:
     elif cfg.model.name == "lt_aware_ce_regression":
         loss_fn = RegressionLossCE(prob_channels=prob_channels)
 
-    # Instantiate the optimizer
-    if encoder_net is not None and cfg.task != "sfm_two_stage":
-        params = list(denoiser_net.parameters()) + list(encoder_net.parameters())
-    else:
-        params = model.parameters()
-
     optimizer = torch.optim.Adam(
-        params=params, lr=cfg.training.hp.lr, betas=[0.9, 0.999], eps=1e-8
+        params=model.parameters(), lr=cfg.training.hp.lr, betas=[0.9, 0.999], eps=1e-8
     )
 
     # Record the current time to measure the duration of subsequent operations.