From c6df63a58817b6414f8a3ae28edd9e6552be3914 Mon Sep 17 00:00:00 2001 From: William Falcon Date: Sun, 4 Oct 2020 21:30:33 -0400 Subject: [PATCH] Fixes #2479 (#3856) --- pytorch_lightning/trainer/__init__.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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 ^^^^^^^^^^^^^^^^