Skip to content

Commit cfe87a0

Browse files
authored
Clarify cluster advanced docs (#16403)
1 parent 60f9406 commit cfe87a0

File tree

1 file changed

+4
-7
lines changed

1 file changed

+4
-7
lines changed

docs/source-pytorch/clouds/cluster_advanced.rst

Lines changed: 4 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -35,21 +35,18 @@ To train a model using multiple nodes, do the following:
3535
.. testcode::
3636

3737
# train.py
38-
def main(hparams):
39-
model = LightningTemplateModel(hparams)
38+
def main(args):
39+
model = YourLightningModule(args)
4040

4141
trainer = Trainer(accelerator="gpu", devices=8, num_nodes=4, strategy="ddp")
4242

4343
trainer.fit(model)
4444

4545

4646
if __name__ == "__main__":
47-
root_dir = os.path.dirname(os.path.realpath(__file__))
48-
parent_parser = ArgumentParser(add_help=False)
49-
hyperparams = parser.parse_args()
50-
47+
args = ... # you can use your CLI parser of choice, or the `LightningCLI`
5148
# TRAIN
52-
main(hyperparams)
49+
main(args)
5350

5451
4. Create the appropriate SLURM job:
5552

0 commit comments

Comments
 (0)