From e2672ee33777ec42540fe0ecb45fc0eee153250b Mon Sep 17 00:00:00 2001
From: John St John <jstjohn@nvidia.com>
Date: Fri, 16 Aug 2024 16:41:09 +0000
Subject: [PATCH] try setting the default root dir in predict to avoid writing
 artifacts to cwd

Signed-off-by: John St John <jstjohn@nvidia.com>
---
 tests/collections/llm/test_mnist_model_nemo2.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tests/collections/llm/test_mnist_model_nemo2.py b/tests/collections/llm/test_mnist_model_nemo2.py
index 96e4aefbf228..a09ed503aa89 100644
--- a/tests/collections/llm/test_mnist_model_nemo2.py
+++ b/tests/collections/llm/test_mnist_model_nemo2.py
@@ -569,6 +569,7 @@ def test_train_mnist_litautoencoder_with_megatron_strategy_single_gpu(tmpdir):
             accelerator="gpu",
             devices=1,
             strategy=pred_strategy,
+            default_root_dir=root_dir,
         )
         ckpt_path = checkpoint_callback.last_model_path.replace(
             ".ckpt", ""