Use NLPDDPStrategyNotebook in Multitask_Prompt_and_PTuning.ipynb

Signed-off-by: Abhishree <[email protected]>
NVIDIA · Dec 20, 2023 · 546eff8 · 546eff8
1 parent 03d9063
commit 546eff8
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/tutorials/nlp/Multitask_Prompt_and_PTuning.ipynb b/tutorials/nlp/Multitask_Prompt_and_PTuning.ipynb
@@ -596,7 +596,7 @@
     "source": [
      "import torch\n",
      "import pytorch_lightning as pl\n",
-     "from nemo.collections.nlp.parts.nlp_overrides import NLPDDPStrategy\n",
+     "from nemo.collections.nlp.parts.nlp_overrides import NLPDDPStrategyNotebook\n",
      "from pytorch_lightning.plugins.environments import TorchElasticEnvironment\n",
      "\n",
      "# let's modify some trainer configs\n",
@@ -618,7 +618,7 @@
      "os.environ[\"RANK\"] = '0'\n",
      "os.environ[\"WORLD_SIZE\"] = '1'\n",
      "\n",
-     "strategy = NLPDDPStrategy(find_unused_parameters=False, no_ddp_communication_hook=True)\n",
+     "strategy = NLPDDPStrategyNotebook(find_unused_parameters=False, no_ddp_communication_hook=True)\n",
      "plugins = [TorchElasticEnvironment()]\n",
      "trainer = pl.Trainer(plugins= plugins, strategy=strategy, **config.trainer)\n",
      "\n",
@@ -783,4 +783,4 @@
   },
   "nbformat": 4,
   "nbformat_minor": 5
- }
+ }