From 3da62ff8da599dfd2b3dfdc263a54b92cc45a868 Mon Sep 17 00:00:00 2001
From: otaj <6065855+otaj@users.noreply.github.com>
Date: Fri, 21 Oct 2022 08:08:11 +0000
Subject: [PATCH] Proper check for availability of bagua (#15220)

---
 src/pytorch_lightning/strategies/bagua.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/pytorch_lightning/strategies/bagua.py b/src/pytorch_lightning/strategies/bagua.py
index 8ec4b29a3e966..8967f73633b78 100644
--- a/src/pytorch_lightning/strategies/bagua.py
+++ b/src/pytorch_lightning/strategies/bagua.py
@@ -16,7 +16,7 @@
 from typing import Any, Dict, List, Optional, Union
 
 import torch
-from lightning_utilities.core.imports import package_available
+from lightning_utilities.core.imports import module_available
 from torch import Tensor
 from torch.nn import Module
 
@@ -32,7 +32,7 @@
 from pytorch_lightning.trainer.states import TrainerFn
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 
-_BAGUA_AVAILABLE = package_available("bagua")
+_BAGUA_AVAILABLE = module_available("bagua.torch_api")
 
 if _BAGUA_AVAILABLE:
     import bagua.torch_api as bagua