From e5310dd982cc85bb5a206f9323f28b82c192a1d2 Mon Sep 17 00:00:00 2001
From: joelamzn <138074614+joelamzn@users.noreply.github.com>
Date: Wed, 3 Apr 2024 14:59:54 -0700
Subject: [PATCH] Import rank_zero_warn

---
 src/neuronx_distributed/lightning/module.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/neuronx_distributed/lightning/module.py b/src/neuronx_distributed/lightning/module.py
index 4c09829..0046cd6 100644
--- a/src/neuronx_distributed/lightning/module.py
+++ b/src/neuronx_distributed/lightning/module.py
@@ -4,6 +4,7 @@
 import torch
 import torch_xla.core.xla_model as xm
 from lightning_utilities.core.apply_func import apply_to_collection
+from lightning_utilities.core.rank_zero import rank_zero_warn
 from pytorch_lightning import LightningModule
 from pytorch_lightning.trainer.connectors.logger_connector.fx_validator import (
     _FxValidator,