Proper check for availability of bagua (#15220)

Lightning-AI · Oct 21, 2022 · 3da62ff · 3da62ff
1 parent 90e1a0e
commit 3da62ff
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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