cleanup and added logic to train for load/save checkpoints (#32)

* added checkpointing logic

* small nits

* check DDP running correctly

* cleanup and added logic to train for load/save checkpoints

* lint

scripts/run_train_nerf.py

@@ -114,7 +114,6 @@ def _train(local_rank: int, world_size: int, config: DictConfig) -> Any:
     """
     trainer = Trainer(config, local_rank, world_size)
     trainer.setup_graph()
-    trainer.setup_optimizer()
     trainer.setup_dataset()
     trainer.train()
     return 0