diff --git a/pytorch_lightning/trainer/__init__.py b/pytorch_lightning/trainer/__init__.py index 8c08d54d9530b..fd25ca129daf8 100644 --- a/pytorch_lightning/trainer/__init__.py +++ b/pytorch_lightning/trainer/__init__.py @@ -358,6 +358,7 @@ def on_train_end(self, trainer, pl_module): cluster_environment ^^^^^^^^^^^^^^^^^^^ Environment to connect arbitrary cluster backends. Lightning automatically handles: + - SLURM - TorchElastic @@ -365,7 +366,7 @@ def on_train_end(self, trainer, pl_module): .. code-block:: python - from pytorch_lightning.cluster_environments import ClusterEnvironment + from pytorch_lightning.cluster_environments import cluster_environment class MyCluster(ClusterEnvironment): @@ -378,6 +379,8 @@ def master_port(self): def world_size(self): return the_world_size + trainer = Trainer(cluster_environment=cluster_environment()) + default_root_dir ^^^^^^^^^^^^^^^^